[python] Bound memory when sampling vectors for index training

paimon-python/pypaimon/globalindex/vindex/vindex_vector_index_writer.py

@@ -106,18 +106,17 @@ def finish(self) -> List[ResultEntry]:
 
             self._close_temp_files()
             self._file_io.check_or_mkdirs(self._index_path)
-            vectors = np.fromfile(
+            training_vectors = _read_training_vectors(
+                np,
                 self._vector_temp_path,
-                dtype=np.float32,
-                count=self._vector_count * self._dimension,
-            ).reshape(self._vector_count, self._dimension)
-            training_vectors = _sample_training_vectors(
-                np, vectors, self._train_sample_ratio)
+                self._vector_count,
+                self._dimension,
+                self._train_sample_ratio,
+            )
             training = VectorIndexTrainer.train(
                 self._training_options(), training_vectors)
             try:
                 del training_vectors
-                del vectors
                 with VectorIndexWriter(training) as writer:
                     self._add_vectors_in_batches(np, writer)
                     with self._file_io.new_output_stream(file_path) as output_stream:
@@ -343,16 +342,43 @@ def _is_float_type(data_type: DataType) -> bool:
     )
 
 
-def _sample_training_vectors(np, vectors, sample_ratio: float):
-    vector_count = vectors.shape[0]
+def _read_training_vectors(
+    np, path: str, vector_count: int, dimension: int, sample_ratio: float
+):
+    """Read the deterministic training sample without loading the full shard.
+
+    Keep the existing evenly spaced sample positions, but gather them from
+    bounded file reads. Only the training sample and a read block need to be
+    resident; blocks containing no selected vectors are skipped entirely.
+    """
     train_count = max(1, min(vector_count, int(math.ceil(
         vector_count * sample_ratio))))
     if train_count == vector_count:
-        return vectors
+        return np.fromfile(
+            path, dtype=np.float32, count=vector_count * dimension,
+        ).reshape(vector_count, dimension)
+
     indexes = (
         np.arange(train_count, dtype=np.int64) * vector_count // train_count
     )
-    return np.ascontiguousarray(vectors[indexes])
+    training_vectors = np.empty((train_count, dimension), dtype=np.float32)
+    item_size = np.dtype(np.float32).itemsize
+    position = 0
+    with open(path, "rb") as vector_file:
+        while position < train_count:
+            start = int(indexes[position])
+            end = min(start + ADD_BATCH_SIZE, vector_count)
+            next_position = int(np.searchsorted(indexes, end))
+            vector_file.seek(start * dimension * item_size)
+            vectors = np.fromfile(
+                vector_file, dtype=np.float32,
+                count=(end - start) * dimension,
+            ).reshape(end - start, dimension)
+            training_vectors[position:next_position] = vectors[
+                indexes[position:next_position] - start]
+            del vectors
+            position = next_position
+    return training_vectors
 
 
 def _materialize_vector(

paimon-python/pypaimon/tests/global_index_build_test.py

@@ -21,7 +21,9 @@
 import os
 import struct
 import sys
+import tempfile
 import types
+from unittest.mock import patch
 
 import pyarrow as pa
 
@@ -41,7 +43,7 @@
 )
 from pypaimon.globalindex.vindex.vindex_vector_index_writer import (
     VindexVectorIndexWriter,
-    _sample_training_vectors,
+    _read_training_vectors,
     native_options,
     train_sample_ratio,
 )
@@ -1027,12 +1029,108 @@ def test_vindex_training_sample_ratio(self):
             0.5, train_sample_ratio(options, 'ivf-rq', 'other'))
 
         import numpy as np
-        vectors = np.arange(20, dtype=np.float32).reshape(10, 2)
-        sampled = _sample_training_vectors(np, vectors, 0.4)
+        with tempfile.TemporaryDirectory() as directory:
+            path = os.path.join(directory, 'vectors.bin')
+            for count, ratio, expected_positions in [
+                (10, 0.4, [0, 2, 5, 7]),
+                (10, 0.31, [0, 2, 5, 7]),
+                (10, 0.01, [0]),
+                (10, 0.99, list(range(10))),
+                (10, 1.0, list(range(10))),
+                (1, 0.01, [0]),
+                (100, 0.03, [0, 33, 66]),
+            ]:
+                with self.subTest(count=count, ratio=ratio):
+                    vectors = np.arange(
+                        count * 2, dtype=np.float32).reshape(count, 2)
+                    vectors.tofile(path)
+                    with patch(
+                        'pypaimon.globalindex.vindex.'
+                        'vindex_vector_index_writer.ADD_BATCH_SIZE', 3
+                    ), patch.object(np, 'fromfile', wraps=np.fromfile) as read:
+                        sampled = _read_training_vectors(
+                            np, path, count, 2, ratio)
+                    np.testing.assert_array_equal(
+                        sampled, vectors[expected_positions])
+                    self.assertEqual(np.float32, sampled.dtype)
+                    self.assertTrue(sampled.flags.c_contiguous)
+                    self.assertTrue(sampled.flags.writeable)
+                    if len(expected_positions) < count:
+                        # A small training sample must not allocate a full
+                        # shard via a single file read.
+                        self.assertTrue(all(
+                            call.kwargs['count'] <= 3 * 2
+                            for call in read.call_args_list))
+                    if count == 100:
+                        self.assertEqual(3, read.call_count)
+
+    def test_vindex_writer_sampled_training_adds_all_vectors(self):
+        schema = pa.schema([
+            ('id', pa.int32()),
+            ('embedding', pa.list_(pa.float32())),
+        ])
+        table = self._create_table(pa_schema=schema, options=self.table_options)
+        writer = VindexVectorIndexWriter(
+            table.file_io,
+            table.path_factory().global_index_path_factory().global_index_root_path(),
+            ArrayType(True, AtomicType('FLOAT')),
+            'ivf-flat',
+            {'ivf-flat.dimension': '2', 'ivf-flat.train.sample-ratio': '0.4'},
+            'embedding',
+        )
+        self.addCleanup(writer.close)
+        writer.write(None, 0)
+        for i in range(10):
+            writer.write([float(i), float(i + 1)], i + 1)
+        paths = [writer._row_id_temp_path, writer._vector_temp_path]
+        _FakeVectorIndexWriter.instances = []
+        with patch.dict(sys.modules, paimon_vindex=types.SimpleNamespace(
+            VectorIndexTrainer=_FakeVectorIndexTrainer,
+            VectorIndexWriter=_FakeVectorIndexWriter,
+        )):
+            entries = writer.finish()
+
+        self.assertEqual(1, len(entries))
+        self.assertEqual(11, entries[0].row_count)
+        built = _FakeVectorIndexWriter.instances[0]
+        self.assertEqual(
+            [[0.0, 1.0], [2.0, 3.0], [5.0, 6.0], [7.0, 8.0]],
+            built.trained)
+        self.assertEqual('10', built.options['expected-vector-count'])
+        self.assertEqual(list(range(1, 11)), built.added_ids)
         self.assertEqual(
-            [[0.0, 1.0], [4.0, 5.0], [10.0, 11.0], [14.0, 15.0]],
-            sampled.tolist(),
+            [[float(i), float(i + 1)] for i in range(10)], built.added_vectors)
+        self.assertTrue(built.closed)
+        self.assertTrue(all(not os.path.exists(path) for path in paths))
+
+    def test_vindex_writer_training_failure_cleans_temp_files(self):
+        schema = pa.schema([
+            ('id', pa.int32()),
+            ('embedding', pa.list_(pa.float32())),
+        ])
+        table = self._create_table(pa_schema=schema, options=self.table_options)
+        writer = VindexVectorIndexWriter(
+            table.file_io,
+            table.path_factory().global_index_path_factory().global_index_root_path(),
+            ArrayType(True, AtomicType('FLOAT')),
+            'ivf-flat',
+            {'ivf-flat.dimension': '2', 'ivf-flat.train.sample-ratio': '0.5'},
+            'embedding',
         )
+        self.addCleanup(writer.close)
+        writer.write([1.0, 0.0], 0)
+        writer.write([0.0, 1.0], 1)
+        paths = [writer._row_id_temp_path, writer._vector_temp_path]
+        with patch.dict(sys.modules, paimon_vindex=types.SimpleNamespace(
+            VectorIndexTrainer=_FakeVectorIndexTrainer,
+            VectorIndexWriter=_FakeVectorIndexWriter,
+        )), patch.object(
+            _FakeVectorIndexTrainer, 'train', side_effect=RuntimeError('training failed')
+        ), self.assertRaisesRegex(RuntimeError, 'training failed'):
+            writer.finish()
+
+        self.assertTrue(all(not os.path.exists(path) for path in paths))
+        self.assertFalse(table.file_io.exists(writer._file_path()))
 
     def test_split_by_contiguous_row_range_matches_java_builder(self):
         split = _FakeSplit([