Skip to content

perf(vindex): optimize IVF index build reads - #800

Merged
JingsongLi merged 25 commits into
apache:mainfrom
jerry-024:codex/ivf-sparse-read-add-clean
Sep 21, 2026
Merged

JingsongLi merged 25 commits into
apache:mainfrom
jerry-024:codex/ivf-sparse-read-add-clean

Conversation

@jerry-024

@jerry-024 jerry-024 commented Sep 10, 2026

Copy link
Copy Markdown
Contributor

Summary

Use a granule-based pipeline by default for IVF index builds. It reads the granules needed for training first, overlaps training with the remaining source reads, and then replays spilled batches and streams live batches into Add. This avoids the previous full-shard raw-vector spill/replay in the normal path.

Implementation

  • Parquet uses page granules when usable OffsetIndex metadata is available and parquet.filter.columnindex.enabled=true; otherwise it uses row groups. Other formats use file granules. Missing or inconsistent metadata safely falls back to a whole-shard granule.
  • Training rows are selected deterministically in logical row space. The sparse-first plan is used only when its estimated read cost is at most 30%; otherwise the pipeline reads the whole shard first.
  • Training, complementary reads, spill replay, and Add overlap through bounded queues, with row range, row ID, and count validation.
  • Long-lived Spill and Add channel workers use named native threads instead of Tokio's shared blocking pool. Training and Replay remain blocking tasks, so the pipeline also completes with max_blocking_threads(1).
  • Predicate-free Parquet RowSelection reads use ordered row-group concurrency and the shared read budget.
  • The pipeline is enabled by default for IVF-Flat, IVF-PQ, IVF-SQ, and IVF-RQ. Set vindex.build.granule.enabled=false to restore the full-spill path. DiskANN is unchanged. There are no public API or on-disk format changes.

Distribution-skew caveat

This issue is limited to the Rust IVF Granule index-build path introduced by this PR; it is not data loss in normal Paimon writes.

The representative risk identified at d1f86cd is a single index build whose source batches have materially different distributions, with a minority category physically concentrated in only a few files or Parquet pages. The Granule training sample can miss or underrepresent that category. All source rows are still validated and passed to Add through spill replay or the live queue, so the index row count remains complete. The possible failure is training coverage: IVF centroids or quantizer codebooks may not model the minority distribution well, which can reduce recall for queries targeting that category at a finite nprobe.

Logical-row stratification, snapshot-derived sampling, and the targeted periodic/unequal/oversized-granule regressions reduce this risk but do not guarantee semantic-category coverage for every physical layout. Set vindex.build.granule.enabled=false when the previous full-spill sampling behavior is required; extremely rare categories still need workload-specific recall validation.

Matched 10M results

One matched trial per index compared latest main f33dee1 with PR head 8af080f. The PR benchmark binary was built from d1f86cd plus the exact blocking-pool patch committed as 8af080f (eed7ac2a141dbc37fd3570ab8c5d78ba49e6b5da56268bbab514ce9c162e0bf1). Binary, source, query-data, and ground-truth checksums were recorded.

Each table contained 10M x 768 Cohere vectors in 10 Parquet files and one index shard. Both indexes used cosine distance and nlist=4096; IVF-PQ additionally used pq.m=192. The host had 8 cores / 16 threads and 61 GiB RAM. Both variants used Rust 1.94.0, registry paimon-vindex-core 0.5.0, Rayon 16, builder concurrency 32, row-group parallelism 8, a 768 MiB read budget, direct OSS, and no Paimon local cache.

Index main build PR build Change Peak RSS, main → PR Recall@10, main → PR Recall@100, main → PR
IVF-PQ 200.566 s 142.238 s -29.1% 4,205.4 → 4,942.4 MiB 0.9440 → 0.9440 0.9170 → 0.9170
IVF-SQ 268.924 s 210.347 s -21.8% 11,802.8 → 11,967.5 MiB 0.9690 → 0.9710 0.9443 → 0.9447

All four builds indexed 10M rows and retained 262,144 training rows. Each main build wrote 28.610 GiB through the old raw tempfile. The PR reported raw_temp_bytes=0; Granule spill was 7.658 GiB for IVF-PQ and 5.035 GiB for IVF-SQ. The PR index files were 4,252 bytes and 2,026 bytes larger, respectively.

Recall used the same 100 queries and ground truth, 10 warmups, nprobe=128, individual mode, request concurrency 1, and one trial. IVF-SQ Recall deltas were +0.0020 and +0.0004; this single trial does not establish a quality improvement. These are single-run comparisons, not variance estimates. Run order was main then PR for each index, and OS page-cache state was not controlled.

Validation

  • Latest CI: all 14 checks passed for 8af080f.
  • The max_blocking_threads(1) regression completes the full Spill → Training → Replay → Add flow; the previous implementation times out.
  • All 18 pipeline unit tests pass.
  • cargo +1.94.0 clippy -p paimon --lib --tests -- -D warnings passes.
  • Matched 10M IVF-PQ and IVF-SQ builds plus Recall@10/Recall@100 completed with identical per-index source files and verified binary/data checksums.
  • Periodic, unequal, and oversized-granule build/commit/query regressions passed across three snapshot-derived samples, with 10/10 target-cluster hits.

@jerry-024 jerry-024 changed the title feat: optimize IVF index build reads perf(vindex): optimize IVF index build reads Sep 10, 2026

@JingsongLi JingsongLi left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Found two reproducible regressions in the sparse training path: sampling can omit an entire data distribution, and the offset-index capability check can enable a second effectively full source read.

Comment on lines +91 to +95
let gap = skipped_rows / gap_count
+ usize::from(
(gap_index + gap_count - gap_extra_offset) % gap_count < skipped_rows % gap_count,
);
cursor = checked_add_offset(cursor, gap, "training gap")?;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Preserve representative training samples across the shard

These fixed gaps permanently exclude contiguous regions; the seed only shifts the remainder allocation by approximately one row. For a 100,000-row shard with the default sample ratio of 1 and 65,536 retained vectors, the last sampled row is 99,468, so a newly appended distribution in the final 500 rows is never trained.

I reproduced this using the actual planner and vindex 0.4 trainer: 90,000 vectors [0], then 9,500 [1], then 500 [100], with IVF-SQ, L2, and nlist=nprobe=1. Querying [100] returns 10/10 results from the final cluster with the previous full-stream reservoir sampling, but 0/10 with these ranges. The SQ bounds are trained only on the older distributions, so the new vectors are clamped to the old upper bound even though every vector is subsequently added and every list is searched. Please retain the existing reservoir path until representative sampling is preserved, and add a recall regression test for data clustered by append order.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verified at 28d6670: both original tail-distribution cases (100,000 rows with a final 500-row cluster, and 1,000,000 rows with a final 10,000-row cluster) now return 10/10 correct-cluster results. The specific fixed-gap regression described above is addressed. A separate small-sample case still fails and is documented on the current planner line here: #800 (comment)

Comment on lines +162 to +168
has_usable_offset_index(
Box::new(input.reader().await?),
file_size,
index_column,
&local_ranges,
)
.await

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Check actual page savings before enabling the extra source pass

An offset index being present does not mean these ranges avoid reading pages. With the default 100,000-row shard, 65,536 training rows, 128-dimensional vectors, and default Zstd/page settings, the 64 sample ranges leave gaps too small to skip pages or survive the reader's existing 1 MiB range coalescing. A tracking FileRead measured exactly 47,788,791 data bytes for both the sample and the full read, excluding metadata; the offset-index check still returned true. This fixture fits the default writer-buffer and file-size limits.

The subsequent full scan therefore doubles source data reads relative to the previous single source scan, including remote reads when the files are on OSS. This measures bytes, not overall wall time versus the saved local spill. Please estimate selected page ranges after coalescing and fall back, or adapt the sample plan, when the sample would read essentially the whole projection.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rechecked at 28d6670: this remains reproducible with the current 512 short, jittered ranges. For the same default 100,000-row shard with 128-dimensional vectors, the sample and full scan still each fetch exactly 47,788,791 data bytes, excluding metadata. The latest concurrency and selected-page budget changes do not prevent the extra effectively full source pass. The sparse-path gate at writer.rs:173 still only checks offset-index availability; please include actual page savings after range coalescing, or fall back when the sample reads essentially the entire projection.

@JingsongLi JingsongLi left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rechecked at 28d6670. The previous 100,000-row and 1,000,000-row tail-distribution cases now pass. Two additional reproducible issues remain in small-sample planning and sparse-read admission; the source-read amplification also remains, with updated measurements in its existing inline thread. Validation: 24 index-build tests and 53 Parquet tests passed, plus four isolated verification cases.

source: None,
});
}
let range_count = training_rows.div_ceil(MAX_IVF_TRAINING_RANGE_ROWS);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Keep small training samples distributed across the shard

When training_rows <= 128, this calculation produces just one contiguous range. A realistic incremental shard with 1,000 new rows and train.sample-ratio=0.1 therefore trains on 100 adjacent rows, whereas the previous implementation sampled every tenth row across the whole shard.

Using the actual planner and vindex 0.4 trainer, I reproduced this with 450 vectors [0], 450 [1], and 100 [100], IVF-SQ, L2, and nlist=nprobe=1. For snapshot 1, bucket 0, and an empty partition, the new range is [385,484], which excludes the entire final distribution. Querying [100] returns 10/10 results from the final cluster with the baseline, but 0/10 with this plan: SQ clamps those vectors to the older upper bound and returns rows 450–459. Please retain multiple strata for small samples, or fall back to the original sampling path, and cover this incremental-shard case in a recall regression test.

Comment on lines +758 to +763
let selected_compressed_bytes = selection
.scan_ranges(page_locations)
.into_iter()
.try_fold(dictionary_bytes, |total, range| {
total.checked_add(range.end.checked_sub(range.start)?)
})?;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Include retained coalesced buffers in sparse read admission

scan_ranges sums the selected pages before ArrowFileReader merges byte ranges separated by at most 1 MiB. Its returned Bytes slices retain the larger coalesced allocations, but the new concurrent row-group admission uses only the smaller selected-page estimate. This leaves live source buffers out of the memory estimate used to increase parallelism.

I verified this with four valid Parquet row groups of 16,384 rows × 128 floats, default Zstd/page settings, and one selected row every 4,096 rows. With a 20,971,520-byte budget, all four groups were admitted and simultaneously retained 25,390,984 bytes of owned read buffers. Allocation ownership and release were tracked with Bytes::from_owner and Drop; these counts exclude decoded Arrow arrays, and all tracked buffers were released on completion. Please account for the coalesced allocations when calculating admission costs, or retain conservative full-column admission when they cannot be estimated.

@JingsongLi JingsongLi left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Optimization can also be targeted at the reading and writing of temporary files. Techniques such as high concurrency and pipelining can be applied to temporary files to address the uncertainties associated with reading remote Parquet files twice.

@JingsongLi JingsongLi left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Requirement fit: SUPPORTED
Implementation: CLEAN

The baseline's full raw-vector spill/replay is a concrete IVF build bottleneck, and the new eligible sparse-read/granule path retains conservative fallbacks for unavailable page metadata, full scans, and DiskANN. I found no P1/P2 regressions in the concurrent Parquet selection path, shared read budget/cancellation handling, sparse-sample row-ID validation, or IVF pipeline completion and cleanup paths.

Validation: cargo test -p paimon --lib vindex_index_build_builder::tests (18 passed) and cargo test -p paimon --lib arrow::format::parquet::tests (54 passed).

@jerry-024
jerry-024 force-pushed the codex/ivf-sparse-read-add-clean branch from 86d753d to d7a207b Compare September 17, 2026 08:36
}

fn pick_indices(total: usize, count: usize) -> impl Iterator<Item = usize> {
(0..count).map(move |index| ((2 * index + 1) * total) / (2 * count))

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Avoid fixed midpoint granule sampling that can exclude an entire distribution

Verified on f934e8a0fdae78d382adaffe2b05696e30070398, using the PR's locked paimon-vindex-core 0.4.0 dependency.

The midpoint selection here is deterministic systematic sampling, rather than randomized stratified sampling. For 4,096 granules and a target of 512, it always selects zero-based granule indices 4, 12, 20, .... Data correlated with file/page order can therefore be completely absent from training. The trainer's subsequent reservoir sampling cannot recover a distribution that never reaches it.

I reproduced this through the actual granule planner, Parquet reads, index build/serialization, and search, not just a simulation of the selection formula:

  • 4,096 real Parquet files in the memory filesystem, 256 rows per file, totaling 1,048,576 rows in one shard.
  • Every eighth file (zero-based file index % 8 == 7) contains [100.0]; the other files alternate [0.0] and [1.0]. The excluded cluster represents 12.5% of all rows.
  • IVF-SQ, L2, dimension=1, nlist=nprobe=1, and train.sample-ratio=1.0.
  • The planner selects 512 first ranges / 131,072 candidate rows, containing no [100.0] vectors.
  • With vindex.build.granule.enabled=false, querying [100.0] returns 10/10 results from that cluster. With it enabled, the same query returns 0/10. These are cluster-hit counts, not exact-ID recall, since distances tie within a cluster.

All rows are still added to the index; the regression is in training quality. SQ learns only the [0, 1] range and clips the unseen [100] vectors to the same upper code. I also reproduced a tail-only distribution failure in a separate actual-planner-plus-trainer comparison.

Could we avoid fixed granule positions and add a regression test for periodic/file-correlated distributions? One conservative option is seeded randomized stratification over logical row positions, expanding selected rows to pages for I/O without treating every fetched row as a training sample. If the required page reads eliminate the savings, retain the full-shard sampling fallback. Randomizing whole granules is less intrusive, but still needs explicit evaluation of rare-cluster coverage and unequal granule sizes. Until that trade-off is validated, I would keep the conservative path as the default.

The existing 24 vindex-build tests and 58 Parquet tests pass; they do not cover this case.

@jerry-024
jerry-024 force-pushed the codex/ivf-sparse-read-add-clean branch from cea04f2 to 1a12958 Compare September 20, 2026 02:01
index_add,
serialize_upload,
rows: row_count_usize,
training_rows_seen: plan.first_rows,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Report the actual trainer sample counts

These fields no longer match the established full-spill semantics. Only training_rows are passed to VectorIndexTrainer, and its reservoir retains the already-computed retained count. The current values therefore overstate the metrics on realistic builds: with a coarse 1,000-row granule and train.sample-ratio=0.1, the trainer sees 100 rows but training_rows_seen reports 1,000; with 10M rows and nlist=4096, the core retains 262,144 rows while training_rows_retained can report 698,768. This makes the default granule-path diagnostics misleading and prevents reliable comparisons with the full-spill path. Please use training_rows_seen: training_rows and training_rows_retained: retained.

@JingsongLi JingsongLi left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1

@JingsongLi
JingsongLi merged commit a48c01b into apache:main Sep 21, 2026
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants