diff --git a/paimon-python/pypaimon/globalindex/create_global_index.py b/paimon-python/pypaimon/globalindex/create_global_index.py index bd9d34a05a7a..b2e125e42159 100644 --- a/paimon-python/pypaimon/globalindex/create_global_index.py +++ b/paimon-python/pypaimon/globalindex/create_global_index.py @@ -21,6 +21,7 @@ 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 @@ 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 @@ def _build_generic_index( 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 compare(left, right): 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 aa56a8ab5975..87c334fa1b2f 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 @@ def write(self, vector, relative_row_id: int) -> None: 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 000000000000..72da3bb74531 --- /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()