[python] Bound memory when sampling vectors for index training - #9750
[python] Bound memory when sampling vectors for index training#9750TheR1sing3un wants to merge 1 commit into
Conversation
JingsongLi
left a comment
There was a problem hiding this comment.
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.
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.