This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 2c8c19c30e [python] Batch vector writes when building global indexes
(#9752)
2c8c19c30e is described below
commit 2c8c19c30e2fc7757207cf56b3fd9051510c7ab8
Author: chaoyang <[email protected]>
AuthorDate: Sun Sep 13 16:24:18 2026 +0800
[python] Batch vector writes when building global indexes (#9752)
---
.../pypaimon/globalindex/create_global_index.py | 38 +++-
.../vindex/vindex_vector_index_writer.py | 66 +++++++
.../pypaimon/tests/vindex_batch_write_test.py | 195 +++++++++++++++++++++
3 files changed, 292 insertions(+), 7 deletions(-)
diff --git a/paimon-python/pypaimon/globalindex/create_global_index.py
b/paimon-python/pypaimon/globalindex/create_global_index.py
index bd9d34a05a..b2e125e421 100644
--- a/paimon-python/pypaimon/globalindex/create_global_index.py
+++ b/paimon-python/pypaimon/globalindex/create_global_index.py
@@ -21,6 +21,7 @@ from functools import cmp_to_key
from typing import Dict, List, Optional, Sequence, Union
import pyarrow as pa
+import pyarrow.compute as pc
from pypaimon.common.options.core_options import CoreOptions
from pypaimon.common.options.options import Options
@@ -51,6 +52,7 @@ from
pypaimon.globalindex.vindex.vindex_vector_global_index_reader import (
VINDEX_IDENTIFIERS,
)
from pypaimon.globalindex.vindex.vindex_vector_index_writer import (
+ ADD_BATCH_SIZE,
VindexVectorIndexWriter,
)
from pypaimon.index.index_file_meta import IndexFileMeta
@@ -302,13 +304,20 @@ class GlobalIndexBuilder:
writer = self._create_generic_index_writer(index_path, index_field)
try:
- for value, row_id in _extract_index_rows(
- table,
- self._index_columns[0],
- SpecialFields.ROW_ID.name,
- index_range,
- ):
- writer.write(value, row_id - index_range.from_)
+ if self._index_type in VINDEX_IDENTIFIERS:
+ if table.column(SpecialFields.ROW_ID.name).null_count:
+ raise ValueError("Cannot build global index because
_ROW_ID is null.")
+ for batch in
table.to_batches(max_chunksize=ADD_BATCH_SIZE):
+ _write_vector_batch(
+ writer, batch, self._index_columns[0], index_range)
+ else:
+ for value, row_id in _extract_index_rows(
+ table,
+ self._index_columns[0],
+ SpecialFields.ROW_ID.name,
+ index_range,
+ ):
+ writer.write(value, row_id - index_range.from_)
index_adds = _to_index_manifest_entries(
self._table,
@@ -444,6 +453,21 @@ def _extract_sorted_rows(
return sorted(rows, key=cmp_to_key(compare))
+def _write_vector_batch(writer, batch, index_column, row_range):
+ row_ids = batch.column(SpecialFields.ROW_ID.name)
+ if row_ids.null_count:
+ raise ValueError("Cannot build global index because _ROW_ID is null.")
+ vectors = batch.column(index_column)
+ selected = pc.and_(
+ pc.greater_equal(row_ids, row_range.from_),
+ pc.less_equal(row_ids, row_range.to),
+ )
+ if not pc.all(selected).as_py():
+ row_ids = pc.filter(row_ids, selected)
+ vectors = pc.filter(vectors, selected)
+ writer.write_batch(vectors, pc.subtract(row_ids, row_range.from_))
+
+
def _extract_index_rows(
table: pa.Table,
index_column: str,
diff --git
a/paimon-python/pypaimon/globalindex/vindex/vindex_vector_index_writer.py
b/paimon-python/pypaimon/globalindex/vindex/vindex_vector_index_writer.py
index aa56a8ab59..87c334fa1b 100644
--- a/paimon-python/pypaimon/globalindex/vindex/vindex_vector_index_writer.py
+++ b/paimon-python/pypaimon/globalindex/vindex/vindex_vector_index_writer.py
@@ -86,6 +86,52 @@ class VindexVectorIndexWriter:
self._vector_temp.write(array("f", materialized).tobytes())
self._vector_count += 1
+ def write_batch(self, vectors, relative_row_ids) -> None:
+ """Write Arrow arrays without materializing valid float32 vectors as
lists.
+
+ Unsupported layouts and invalid vectors use the scalar path so that
+ validation errors and the order of successfully written rows match
write().
+ """
+ if self._closed:
+ raise RuntimeError("VindexVectorIndexWriter is already closed.")
+ if len(vectors) != len(relative_row_ids):
+ raise ValueError("Vector and row ID batch lengths differ.")
+ if relative_row_ids.null_count:
+ raise ValueError("Cannot build global index because _ROW_ID is
null.")
+ if len(vectors) == 0:
+ return
+
+ import numpy as np
+ import pyarrow as pa
+ import pyarrow.compute as pc
+
+ valid_vectors, valid_ids = vectors, relative_row_ids
+ if vectors.null_count:
+ valid = pc.is_valid(vectors)
+ valid_vectors = pc.filter(vectors, valid)
+ valid_ids = pc.filter(relative_row_ids, valid)
+ if len(valid_vectors) == 0:
+ self._row_count += len(vectors)
+ return
+
+ values = _float32_batch_values(np, pa, valid_vectors, self._dimension)
+ if (values is not None and values.null_count == 0
+ and valid_ids.type == pa.int64()):
+ data = values.to_numpy(zero_copy_only=True)
+ if np.isfinite(data).all():
+ ids = np.ascontiguousarray(
+ valid_ids.to_numpy(zero_copy_only=False), dtype=np.int64)
+ data = np.ascontiguousarray(data, dtype=np.float32)
+ self._row_count += len(vectors)
+ self._ensure_temp_files()
+ self._row_id_temp.write(memoryview(ids).cast('B'))
+ self._vector_temp.write(memoryview(data).cast('B'))
+ self._vector_count += len(valid_vectors)
+ return
+
+ for vector, row_id in zip(vectors.to_pylist(),
relative_row_ids.to_pylist()):
+ self.write(vector, row_id)
+
def finish(self) -> List[ResultEntry]:
if self._closed:
raise RuntimeError("VindexVectorIndexWriter is already closed.")
@@ -355,6 +401,26 @@ def _sample_training_vectors(np, vectors, sample_ratio:
float):
return np.ascontiguousarray(vectors[indexes])
+def _float32_batch_values(np, pa, vectors, dimension):
+ vector_type = vectors.type
+ if not (
+ pa.types.is_list(vector_type)
+ or pa.types.is_large_list(vector_type)
+ or pa.types.is_fixed_size_list(vector_type)
+ ) or vector_type.value_type != pa.float32():
+ return None
+
+ if pa.types.is_fixed_size_list(vector_type):
+ if vector_type.list_size != dimension:
+ return None
+ return vectors.values.slice(vectors.offset * dimension, len(vectors) *
dimension)
+
+ offsets = vectors.offsets.to_numpy(zero_copy_only=True)
+ if not np.all(np.diff(offsets) == dimension):
+ return None
+ return vectors.values.slice(int(offsets[0]), int(offsets[-1] - offsets[0]))
+
+
def _materialize_vector(
value, dimension: int, relative_row_id: int
) -> List[float]:
diff --git a/paimon-python/pypaimon/tests/vindex_batch_write_test.py
b/paimon-python/pypaimon/tests/vindex_batch_write_test.py
new file mode 100644
index 0000000000..72da3bb745
--- /dev/null
+++ b/paimon-python/pypaimon/tests/vindex_batch_write_test.py
@@ -0,0 +1,195 @@
+# 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 os
+import unittest
+from unittest.mock import Mock, patch
+
+import numpy as np
+import pyarrow as pa
+
+from pypaimon.globalindex.create_global_index import GlobalIndexBuilder,
_write_vector_batch
+from pypaimon.globalindex.vindex.vindex_vector_index_writer import
VindexVectorIndexWriter
+from pypaimon.schema.data_types import ArrayType, AtomicType
+from pypaimon.utils.range import Range
+
+
+class VindexBatchWriteTest(unittest.TestCase):
+
+ def _writer(self):
+ writer = VindexVectorIndexWriter(
+ Mock(), '/unused', ArrayType(True, AtomicType('FLOAT')),
+ 'ivf-flat', {'ivf-flat.dimension': '2'}, 'embedding')
+ self.addCleanup(writer.close)
+ return writer
+
+ def _contents(self, writer):
+ writer._close_temp_files()
+ contents = []
+ for path in (writer._row_id_temp_path, writer._vector_temp_path):
+ if path is None:
+ contents.append(b'')
+ else:
+ with open(path, 'rb') as stream:
+ contents.append(stream.read())
+ return writer._row_count, writer._vector_count, contents
+
+ def _assert_parity(self, vectors, row_ids, fast=True):
+ scalar, batch = self._writer(), self._writer()
+ for vector, row_id in zip(vectors.to_pylist(), row_ids.to_pylist()):
+ scalar.write(vector, row_id)
+ with patch.object(batch, 'write', wraps=batch.write) as write:
+ batch.write_batch(vectors, row_ids)
+ if fast:
+ write.assert_not_called()
+ self.assertEqual(self._contents(scalar), self._contents(batch))
+
+ def test_sliced_list_large_list_and_fixed_size_list(self):
+ for array_type in (pa.list_(pa.float32()), pa.large_list(pa.float32()),
+ pa.list_(pa.float32(), 2)):
+ with self.subTest(array_type=array_type):
+ vectors = pa.array([
+ [99, 99], [1, 2], None, [-0.0, 1e-40], [3, 4], [88, 88],
+ ], type=array_type).slice(1, 4)
+ row_ids = pa.array([77, 0, 3, 8, 2 ** 62, 99],
type=pa.int64()).slice(1, 4)
+ self._assert_parity(vectors, row_ids)
+
+ def test_null_parent_ignores_invalid_child_values(self):
+ values = pa.array([1, 2, float('nan'), None, 3, 4], type=pa.float32())
+ vectors = pa.Array.from_buffers(pa.list_(pa.float32()), 3, [
+ pa.py_buffer(b'\x05'),
+ pa.py_buffer(np.array([0, 2, 4, 6], dtype=np.int32)),
+ ], children=[values])
+ self._assert_parity(vectors, pa.array([0, 1, 2], type=pa.int64()))
+
+ def test_empty_and_all_null_batches_do_not_create_files(self):
+ for data in ([], [None, None]):
+ with self.subTest(data=data):
+ writer = self._writer()
+ writer.write_batch(pa.array(data, type=pa.list_(pa.float32())),
+ pa.array(range(len(data)), type=pa.int64()))
+ self.assertEqual((len(data), 0, [b'', b'']),
self._contents(writer))
+
+ def test_multiple_batches_and_scalar_writes_can_be_interleaved(self):
+ writer = self._writer()
+ writer.write([1, 2], 0)
+ writer.write_batch(pa.array([[3, 4], None],
type=pa.list_(pa.float32())),
+ pa.array([2, 3], type=pa.int64()))
+ writer.write([5, 6], 4)
+ writer.write_batch(pa.array([[7, 8]], type=pa.list_(pa.float32(), 2)),
+ pa.array([6], type=pa.int64()))
+ count, valid_count, contents = self._contents(writer)
+ self.assertEqual((5, 4), (count, valid_count))
+ self.assertEqual([0, 2, 4, 6], np.frombuffer(contents[0],
dtype=np.int64).tolist())
+ self.assertEqual(list(range(1, 9)), np.frombuffer(contents[1],
dtype=np.float32).tolist())
+
+ def test_invalid_vectors_preserve_scalar_error_and_written_prefix(self):
+ cases = [
+ [[1, 2], [3]],
+ [[1, 2], [None, 4]],
+ [[1, 2], [float('nan'), 4]],
+ [[1, 2], [3, float('inf')]],
+ [[1, 2], [3, float('-inf')]],
+ [[float('nan'), 2], [3]],
+ [[1, 2], [float('nan'), None]],
+ [[1, 2], [None, float('nan')]],
+ ]
+ for data in cases:
+ with self.subTest(data=data):
+ vectors = pa.array(data, type=pa.list_(pa.float32()))
+ ids = pa.array([5, 9], type=pa.int64())
+ scalar, batch = self._writer(), self._writer()
+ with self.assertRaises(ValueError) as old_error:
+ for vector, row_id in zip(vectors.to_pylist(),
ids.to_pylist()):
+ scalar.write(vector, row_id)
+ with self.assertRaises(ValueError) as new_error:
+ batch.write_batch(vectors, ids)
+ self.assertEqual(str(old_error.exception),
str(new_error.exception))
+ self.assertEqual(self._contents(scalar), self._contents(batch))
+
+ def test_float64_and_non_int64_ids_fall_back_to_scalar(self):
+ self._assert_parity(
+ pa.array([[1.1, 2.2], None], type=pa.list_(pa.float64())),
+ pa.array([0, 1], type=pa.int64()), fast=False)
+ self._assert_parity(
+ pa.array([[1, 2], [3, 4]], type=pa.list_(pa.float32())),
+ pa.array([0, 1], type=pa.int32()), fast=False)
+
+ def test_batch_length_and_null_row_id_validation(self):
+ vectors = pa.array([[1, 2]], type=pa.list_(pa.float32()))
+ writer = self._writer()
+ with self.assertRaisesRegex(ValueError, 'batch lengths differ'):
+ writer.write_batch(vectors, pa.array([], type=pa.int64()))
+ with self.assertRaisesRegex(ValueError, '_ROW_ID is null'):
+ writer.write_batch(vectors, pa.array([None], type=pa.int64()))
+ self.assertEqual((0, 0, [b'', b'']), self._contents(writer))
+
+ def test_close_removes_batch_files_and_rejects_further_writes(self):
+ writer = self._writer()
+ vectors = pa.array([[1, 2]], type=pa.list_(pa.float32()))
+ ids = pa.array([0], type=pa.int64())
+ writer.write_batch(vectors, ids)
+ paths = writer._row_id_temp_path, writer._vector_temp_path
+ writer.close()
+ self.assertTrue(all(not os.path.exists(path) for path in paths))
+ with self.assertRaisesRegex(RuntimeError, 'already closed'):
+ writer.write_batch(vectors, ids)
+
+ def test_builder_filters_ranges_before_vector_validation(self):
+ vectors = pa.array([[float('nan'), 0], [1, 2], None, [3, 4], [5]],
+ type=pa.list_(pa.float32()))
+ ids = pa.array([9, 10, 11, 19, 20], type=pa.int64())
+ batch = pa.RecordBatch.from_arrays([vectors, ids], ['embedding',
'_ROW_ID'])
+ writer = self._writer()
+ _write_vector_batch(writer, batch, 'embedding', Range(10, 19))
+ count, valid_count, contents = self._contents(writer)
+ self.assertEqual((3, 2), (count, valid_count))
+ self.assertEqual([0, 9], np.frombuffer(contents[0],
dtype=np.int64).tolist())
+ self.assertEqual([1, 2, 3, 4], np.frombuffer(contents[1],
dtype=np.float32).tolist())
+ batch = pa.RecordBatch.from_arrays([
+ pa.array([[1, 2]], type=pa.list_(pa.float32())),
+ pa.array([None], type=pa.int64()),
+ ], ['embedding', '_ROW_ID'])
+ with self.assertRaisesRegex(ValueError, '_ROW_ID is null'):
+ _write_vector_batch(self._writer(), batch, 'embedding', Range(10,
19))
+
+ def test_null_row_ids_are_rejected_before_writing_any_batch(self):
+ builder = object.__new__(GlobalIndexBuilder)
+ builder._core_options = Mock()
+ builder._core_options.global_index_row_count_per_shard.return_value =
10
+ builder._index_type = 'ivf-flat'
+ builder._index_columns = ['embedding']
+ writer = Mock()
+ builder._create_generic_index_writer = Mock(return_value=writer)
+ read = Mock()
+ read.to_arrow.return_value = pa.table({
+ 'embedding': pa.array([[1], [3, 4]], type=pa.list_(pa.float32())),
+ '_ROW_ID': pa.array([0, None], type=pa.int64()),
+ })
+ module = 'pypaimon.globalindex.create_global_index'
+ with patch(module + '._split_by_global_index_shard', return_value=[
+ (Mock(), Range(0, 9)),
+ ]), patch(module + '.ADD_BATCH_SIZE', 1):
+ with self.assertRaisesRegex(ValueError, '_ROW_ID is null'):
+ builder._build_generic_index([], [], Mock(), read, '/unused')
+ writer.write_batch.assert_not_called()
+ writer.finish.assert_not_called()
+ writer.close.assert_called_once()
+
+
+if __name__ == '__main__':
+ unittest.main()
