[python] Stream training samples into native vector index trainers

paimon-python/README.md

@@ -302,3 +302,13 @@ only support `seek` and `read` remain serialized. Workers are created lazily
 and released when the index reader closes; separate readers have separate
 budgets. This option controls index I/O, not shard search or native compute
 threads.
+
+
+# Native vector index training
+
+The native vector index writer submits training vectors in bounded batches.
+`<index-type>.train.sample-ratio` (or its field-level override) still selects
+the same evenly spaced non-null vectors in the same order. Native training
+receives the final corpus size for automatic IVF sizing. This bounds Python
+training buffers; native training and index construction have their own
+memory requirements.

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

@@ -152,18 +152,8 @@ def finish(self) -> List[ResultEntry]:
 
             self._close_temp_files()
             self._file_io.check_or_mkdirs(self._index_path)
-            vectors = np.fromfile(
-                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)
-            training = VectorIndexTrainer.train(
-                self._training_options(), training_vectors)
+            training = self._train(np, VectorIndexTrainer)
             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:
@@ -178,6 +168,16 @@ def finish(self) -> List[ResultEntry]:
 
         return [ResultEntry(self.file_name, self._row_count, b"{}")]
 
+    def _train(self, np, trainer_type):
+        with open(self._vector_temp_path, "rb") as vector_file:
+            with trainer_type.create(self._training_options()) as trainer:
+                for batch in _iter_training_batches(
+                    np, vector_file, self._vector_count, self._dimension,
+                    self._train_sample_ratio, batch_size=ADD_BATCH_SIZE,
+                ):
+                    trainer.add_training_vectors(batch)
+                return trainer.finish_training()
+
     def _file_path(self) -> str:
         return "%s/%s" % (self._index_path, self.file_name)
 
@@ -389,16 +389,31 @@ def _is_float_type(data_type: DataType) -> bool:
     )
 
 
-def _sample_training_vectors(np, vectors, sample_ratio: float):
-    vector_count = vectors.shape[0]
+def _iter_training_batches(
+    np, vector_file, vector_count: int, dimension: int, sample_ratio: float,
+    batch_size: int = ADD_BATCH_SIZE,
+):
+    """Yield the existing evenly spaced sample using bounded reads and buffers."""
     train_count = max(1, min(vector_count, int(math.ceil(
         vector_count * sample_ratio))))
-    if train_count == vector_count:
-        return vectors
-    indexes = (
-        np.arange(train_count, dtype=np.int64) * vector_count // train_count
-    )
-    return np.ascontiguousarray(vectors[indexes])
+    position = 0
+    item_size = np.dtype(np.float32).itemsize
+    while position < train_count:
+        start = position * vector_count // train_count
+        end = min(start + batch_size, vector_count)
+        # First sample position whose source row is at or beyond this block.
+        next_position = min(train_count, (end * train_count + vector_count - 1) // vector_count)
+        vector_file.seek(start * dimension * item_size)
+        vectors = np.fromfile(
+            vector_file, dtype=np.float32, count=(end - start) * dimension,
+        ).reshape(end - start, dimension)
+        if train_count == vector_count:
+            yield vectors
+        else:
+            indexes = np.arange(position, next_position, dtype=np.int64)
+            indexes = indexes * vector_count // train_count - start
+            yield np.ascontiguousarray(vectors[indexes])
+        position = next_position
 
 
 def _float32_batch_values(np, pa, vectors, dimension):

paimon-python/pypaimon/tests/global_index_build_test.py

@@ -22,6 +22,7 @@
 import struct
 import sys
 import types
+import tempfile
 
 import pyarrow as pa
 
@@ -41,7 +42,7 @@
 )
 from pypaimon.globalindex.vindex.vindex_vector_index_writer import (
     VindexVectorIndexWriter,
-    _sample_training_vectors,
+    _iter_training_batches,
     native_options,
     train_sample_ratio,
 )
@@ -93,9 +94,26 @@ def close(self):
 
 class _FakeVectorIndexTrainer:
 
+    def __init__(self, options):
+        self.options = options
+        self.batches = []
+
     @classmethod
-    def train(cls, options, data):
-        return _FakeVectorIndexTraining(options, data)
+    def create(cls, options):
+        return cls(options)
+
+    def add_training_vectors(self, data):
+        self.batches.append(data.copy())
+
+    def finish_training(self):
+        import numpy as np
+        return _FakeVectorIndexTraining(self.options, np.concatenate(self.batches))
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
 
 
 class _FakeVectorIndexWriter:
@@ -1028,7 +1046,11 @@ def test_vindex_training_sample_ratio(self):
 
         import numpy as np
         vectors = np.arange(20, dtype=np.float32).reshape(10, 2)
-        sampled = _sample_training_vectors(np, vectors, 0.4)
+        with tempfile.TemporaryFile() as vector_file:
+            vectors.tofile(vector_file)
+            vector_file.flush()
+            sampled = np.concatenate(list(_iter_training_batches(
+                np, vector_file, 10, 2, 0.4, batch_size=3)))
         self.assertEqual(
             [[0.0, 1.0], [4.0, 5.0], [10.0, 11.0], [14.0, 15.0]],
             sampled.tolist(),

paimon-python/pypaimon/tests/vindex_training_test.py

@@ -0,0 +1,115 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import math
+import os
+import tempfile
+import unittest
+from unittest import mock
+
+import numpy as np
+
+from pypaimon.filesystem.local_file_io import LocalFileIO
+from pypaimon.globalindex.vindex.vindex_vector_index_writer import (
+    VindexVectorIndexWriter, _iter_training_batches,
+)
+from pypaimon.schema.data_types import ArrayType, AtomicType
+
+
+class VindexTrainingTest(unittest.TestCase):
+
+    def test_batches_preserve_sample_positions_and_bound_reads(self):
+        vectors = np.arange(10003 * 3, dtype=np.float32).reshape(-1, 3)
+        with tempfile.TemporaryFile() as stream:
+            vectors.tofile(stream)
+            stream.flush()
+            for ratio in (1.0, 0.999, 0.37, 0.01, 1e-8):
+                for batch_size in (1, 17, 1000):
+                    with self.subTest(ratio=ratio, batch_size=batch_size):
+                        with mock.patch.object(np, "fromfile", wraps=np.fromfile) as read:
+                            batches = list(_iter_training_batches(
+                                np, stream, len(vectors), 3, ratio, batch_size))
+                        count = max(1, math.ceil(len(vectors) * ratio))
+                        indexes = np.arange(count) * len(vectors) // count
+                        np.testing.assert_array_equal(vectors[indexes], np.concatenate(batches))
+                        self.assertTrue(all(b.flags.c_contiguous for b in batches))
+                        self.assertTrue(all(len(b) <= batch_size for b in batches))
+                        self.assertTrue(all(c[1]["count"] <= batch_size * 3
+                                            for c in read.call_args_list))
+
+    def test_training_failure_closes_trainer_and_removes_temp_files(self):
+        with tempfile.TemporaryDirectory() as directory:
+            writer = self._writer(directory, {})
+            writer.write([1.0] * 8, 0)
+            paths = [writer._vector_temp_path, writer._row_id_temp_path]
+            for phase in ("add_training_vectors", "finish_training"):
+                trainer = mock.MagicMock()
+                trainer.__enter__.return_value = trainer
+                getattr(trainer, phase).side_effect = RuntimeError("training failed")
+                module = mock.Mock()
+                module.VectorIndexTrainer.create.return_value = trainer
+                with mock.patch.dict("sys.modules", {"paimon_vindex": module}):
+                    with self.assertRaisesRegex(RuntimeError, "training failed"):
+                        writer.finish()
+                trainer.__exit__.assert_called_once()
+                self.assertTrue(all(not os.path.exists(path) for path in paths))
+                self.assertFalse(os.path.exists(writer._file_path()))
+                writer = self._writer(directory, {})
+                writer.write([1.0] * 8, 0)
+                paths = [writer._vector_temp_path, writer._row_id_temp_path]
+            writer.close()
+
+    def test_native_streamed_build_matches_one_shot(self):
+        try:
+            from paimon_vindex import VectorIndexTrainer, VectorIndexWriter
+        except ImportError:
+            self.skipTest("paimon-vindex is not installed")
+        vectors = np.random.default_rng(42).standard_normal((2049, 8)).astype(np.float32)
+        with tempfile.TemporaryDirectory() as directory:
+            for ratio in (1.0, 0.37):
+                for index_type in ("ivf-flat", "ivf-pq", "ivf-sq", "ivf-rq", "diskann"):
+                    with self.subTest(ratio=ratio, index_type=index_type):
+                        options = {index_type + ".train.sample-ratio": str(ratio)}
+                        if index_type != "diskann":
+                            options[index_type + ".nlist"] = "16"
+                        writer = self._writer(directory, options, index_type)
+                        writer.write(None, 0)
+                        for i, vector in enumerate(vectors):
+                            writer.write(vector, i + 1)
+                        reference_path = os.path.join(directory, "reference")
+                        count = math.ceil(len(vectors) * ratio)
+                        sample = vectors[np.arange(count) * len(vectors) // count]
+                        with VectorIndexTrainer.train(writer._training_options(), sample) as training:
+                            with VectorIndexWriter(training) as native:
+                                native.add_vectors(np.arange(1, len(vectors) + 1), vectors)
+                                with open(reference_path, "wb") as output:
+                                    native.write(output)
+                        with mock.patch(
+                            "pypaimon.globalindex.vindex.vindex_vector_index_writer.ADD_BATCH_SIZE", 127
+                        ):
+                            result = writer.finish()
+                        self.assertEqual(1, len(result))
+                        with open(reference_path, "rb") as reference, open(writer._file_path(), "rb") as actual:
+                            self.assertEqual(reference.read(), actual.read())
+
+    @staticmethod
+    def _writer(directory, options, index_type="ivf-flat"):
+        options = dict(options)
+        options[index_type + ".dimension"] = "8"
+        return VindexVectorIndexWriter(
+            LocalFileIO(), directory, ArrayType(True, AtomicType("FLOAT")),
+            index_type, options, "embedding")