Skip to content

[python] Bound memory when sampling vectors for index training - #9750

Closed
TheR1sing3un wants to merge 1 commit into
apache:masterfrom
TheR1sing3un:codex/vindex-training-sample-memory
Closed

[python] Bound memory when sampling vectors for index training#9750
TheR1sing3un wants to merge 1 commit into
apache:masterfrom
TheR1sing3un:codex/vindex-training-sample-memory

Conversation

@TheR1sing3un

Copy link
Copy Markdown
Member

Purpose

Building a vindex index with a training sample previously loaded the entire temporary vector file into a NumPy array before selecting training rows. A small sample therefore still required memory proportional to the full index shard.

Read selected training vectors in blocks of at most 10,000 rows instead. Preserve the existing sample count, evenly spaced sample positions, float32 layout, and full-vector indexing after training. Skip gaps between read blocks when no training vectors are needed. The full-sample path still reads the complete training array required by the native API.

This change only bounds the sampling-stage read buffers. It does not bound upstream Arrow shard materialization or the native trainer's own allocations.

Tests

  • python -m pytest pypaimon/tests/global_index_build_test.py -k vindex -q: 10 passed, 18 deselected. Covers sample parity, bounded reads, sparse samples, null-vector exclusion, full-vector indexing, and temporary-file cleanup.
  • python -m flake8 --config dev/cfg.ini pypaimon/globalindex/vindex/vindex_vector_index_writer.py pypaimon/tests/global_index_build_test.py: passed.
  • Native paimon-vindex 0.4.0 smoke comparison: 257 eight-dimensional vectors, IVF-Flat with four lists and 25% training sampling. Eight queries with nprobe=4 produced identical top-5 IDs and distances for old and new sampling paths.
  • Isolated sampling memory comparison, in separate macOS processes: a 256 MiB file containing 262,144 256-dimensional float32 vectors, sampled at 1%, produced identical sample SHA-256 hashes. Peak process RSS was 279.64 MiB before and 33.41 MiB after. This is a sampling-only measurement, not an end-to-end index build benchmark; filesystem cache was not controlled.

@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.

Reviewed d65f577. The sampling memory problem is real, and the focused vindex tests pass (10 tests). However, #9758 now covers the same deterministic bounded sampler and goes further by streaming those batches into the native trainer, including the full-sample case. Both PRs replace the same baseline training path; this is an alternative implementation, not an independent step that needs to be merged as well.

I compared this sampler against #9758 across dense, sparse and full sampling ratios, including batch boundaries: all samples matched exactly. #9758 also passed native byte-for-byte index comparisons for the five supported index types. Keeping both would leave duplicate/conflicting sampling implementations and tests without an additional end-to-end capability.

Closing this PR as superseded by #9758; please keep the bounded-sampling coverage in that PR. If this version has a distinct supported use case that the streaming trainer cannot cover, we can reconsider it with that concrete case.

@JingsongLi JingsongLi closed this Sep 12, 2026
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.

2 participants