This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit cc21db42e8c5cabe233547375c9a0c812029b6f7 Author: chaoyang <[email protected]> AuthorDate: Thu Sep 17 17:37:06 2026 +0800 [python] Build global index shards concurrently (#9798) --- .../global-index/manage-indexes.mdx | 1 + .../pypaimon/common/options/core_options.py | 14 ++ .../pypaimon/globalindex/create_global_index.py | 214 +++++++++++++++------ .../pypaimon/tests/global_index_build_test.py | 195 +++++++++++++++++++ .../pypaimon/tests/vindex_batch_write_test.py | 2 + 5 files changed, 369 insertions(+), 57 deletions(-) diff --git a/docs/docs/multimodal-table/global-index/manage-indexes.mdx b/docs/docs/multimodal-table/global-index/manage-indexes.mdx index 2e1ff0c9f9..24f2e70a54 100644 --- a/docs/docs/multimodal-table/global-index/manage-indexes.mdx +++ b/docs/docs/multimodal-table/global-index/manage-indexes.mdx @@ -454,6 +454,7 @@ These table options affect global index build and read behavior: | `sorted-index.records-per-range` | `10000000` | Expected number of records per sorted global index file for BTree, Bitmap, and Multivalue builds. | | `sorted-index.build.max-parallelism` | `4096` | Maximum Flink or Spark parallelism for building sorted global indexes. | | `global-index.row-count-per-shard` | `100000` | Target row count per shard for non-sorted global index builds such as vector and full-text indexes. | +| `global-index.build.parallelism` | `1` | Number of shards built concurrently by the local PyPaimon builder. Each shard may use native worker threads, so increase this value conservatively. | | `global-index.build.max-shard` | `32` | Preferred maximum shard count for global index builds. | | `global-index.build.max-parallelism` | `4096` | Maximum Flink or Spark parallelism for building non-sorted global indexes. | | `global-index.thread-num` | `32` | Maximum number of concurrent threads for global index I/O. | diff --git a/paimon-python/pypaimon/common/options/core_options.py b/paimon-python/pypaimon/common/options/core_options.py index ef9d45fd82..0de00abdc1 100644 --- a/paimon-python/pypaimon/common/options/core_options.py +++ b/paimon-python/pypaimon/common/options/core_options.py @@ -952,6 +952,17 @@ class CoreOptions: .with_description("Row count per shard for global index.") ) + GLOBAL_INDEX_BUILD_PARALLELISM: ConfigOption[int] = ( + ConfigOptions.key("global-index.build.parallelism") + .int_type() + .default_value(1) + .with_description( + "Number of global index shards built concurrently by the local " + "Python builder. Each shard may also use native worker threads, " + "so increase this value conservatively." + ) + ) + PK_VECTOR_INDEX_COLUMNS: ConfigOption[str] = ( ConfigOptions.key("pk-vector.index.columns") .string_type() @@ -1698,6 +1709,9 @@ class CoreOptions: def global_index_row_count_per_shard(self) -> int: return self.options.get(CoreOptions.GLOBAL_INDEX_ROW_COUNT_PER_SHARD) + def global_index_build_parallelism(self) -> int: + return self.options.get(CoreOptions.GLOBAL_INDEX_BUILD_PARALLELISM) + def primary_key_btree_index_columns(self) -> List[str]: return self._primary_key_index_columns(CoreOptions.PK_BTREE_INDEX_COLUMNS) diff --git a/paimon-python/pypaimon/globalindex/create_global_index.py b/paimon-python/pypaimon/globalindex/create_global_index.py index d68221a0ae..93473bb7b3 100644 --- a/paimon-python/pypaimon/globalindex/create_global_index.py +++ b/paimon-python/pypaimon/globalindex/create_global_index.py @@ -17,6 +17,8 @@ """Build global index files from Python.""" +import threading +from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait from functools import cmp_to_key from typing import Dict, List, Optional, Sequence, Union @@ -288,72 +290,170 @@ class GlobalIndexBuilder: def _build_generic_index( self, splits, unindexed_ranges, index_field, table_read, index_path: str ) -> List[CommitMessage]: - from pypaimon.read.table_read import _ClosableArrowBatchReader - rows_per_shard = self._core_options.global_index_row_count_per_shard() if rows_per_shard <= 0: raise ValueError( "Option 'global-index.row-count-per-shard' must be greater than 0." ) - messages = [] - for index_split, index_range in _split_by_global_index_shard( - splits, rows_per_shard, unindexed_ranges - ): - writer = None + parallelism = self._core_options.global_index_build_parallelism() + if parallelism <= 0: + raise ValueError( + "Option 'global-index.build.parallelism' must be greater than 0." + ) + + shards = _split_by_global_index_shard( + splits, rows_per_shard, unindexed_ranges) + if not shards: + return [] + + if parallelism == 1 or len(shards) == 1: + messages = [] try: - reader, batches = table_read._new_arrow_batch_reader([index_split]) - # Close the Python iterator explicitly on failure as well as - # the Arrow reader, which may retain a suspended generator. - with _ClosableArrowBatchReader(reader, batches) as batch_reader: - for batch in batch_reader: - if batch.num_rows == 0: - continue - if writer is None: - writer = self._create_generic_index_writer( - index_path, index_field) - if self._index_type in VINDEX_IDENTIFIERS: - if batch.column(SpecialFields.ROW_ID.name).null_count: - raise ValueError( - "Cannot build global index because _ROW_ID is null.") - for offset in range(0, batch.num_rows, ADD_BATCH_SIZE): - _write_vector_batch( - writer, batch.slice(offset, ADD_BATCH_SIZE), - self._index_columns[0], index_range) - else: - for value, row_id in _extract_index_rows( - batch, - self._index_columns[0], - SpecialFields.ROW_ID.name, - index_range, - ): - writer.write(value, row_id - index_range.from_) - del batch - - if writer is None: - continue - - index_adds = _to_index_manifest_entries( - self._table, - index_split.partition, - index_range, - index_field.id, - self._index_type, - writer.finish(), + for index_split, index_range in shards: + message = self._build_generic_shard( + index_split, index_range, index_field, table_read, index_path) + if message is not None: + messages.append(message) + return messages + except BaseException: + self._delete_uncommitted_indexes(messages) + raise + + # Workers record their own output so that rollback never depends on the + # future list being fully built. ThreadPoolExecutor.submit() enqueues the + # work item before it starts an extra worker, so a submission that raises + # (RuntimeError: can't start new thread) may still run its shard on an + # already running worker. + completed = [] + completed_lock = threading.Lock() + + def build_shard(index_split, index_range): + message = self._build_generic_shard( + index_split, index_range, index_field, table_read, index_path) + if message is not None: + with completed_lock: + completed.append(message) + return message + + futures = [] + try: + with ThreadPoolExecutor( + max_workers=min(parallelism, len(shards)), + thread_name_prefix="paimon-global-index-build", + ) as executor: + try: + for index_split, index_range in shards: + futures.append( + executor.submit(build_shard, index_split, index_range)) + except BaseException: + # Keep queued shards that have not started from building an + # index file this build is about to delete. + for future in futures: + future.cancel() + raise + done, _ = wait(futures, return_when=FIRST_EXCEPTION) + failed = next( + (future for future in futures + if future in done and future.exception() is not None), + None, ) - finally: - if writer is not None: - writer.close() - if index_adds: - messages.append( - CommitMessage( - partition=tuple(index_split.partition.values), - bucket=0, - new_files=[], - index_adds=index_adds, - ) + if failed is not None: + for future in futures: + future.cancel() + failed.result() + + # Futures are consumed in shard-plan order so index manifest + # messages are deterministic even when shards finish out of order. + results = [future.result() for future in futures] + return [message for message in results if message is not None] + except BaseException: + # Exiting the executor waits for in-flight shards to close their + # readers and writers, so ``completed`` is stable and fully visible + # here. Delete every index that was built because build() will not + # return commit messages after a failure. + self._delete_uncommitted_indexes(completed) + raise + + def _build_generic_shard( + self, index_split, index_range, index_field, table_read, index_path: str + ) -> Optional[CommitMessage]: + from pypaimon.read.table_read import _ClosableArrowBatchReader + + writer = None + try: + reader, batches = table_read._new_arrow_batch_reader([index_split]) + # Close the Python iterator explicitly on failure as well as + # the Arrow reader, which may retain a suspended generator. + with _ClosableArrowBatchReader(reader, batches) as batch_reader: + for batch in batch_reader: + if batch.num_rows == 0: + continue + if writer is None: + writer = self._create_generic_index_writer( + index_path, index_field) + if self._index_type in VINDEX_IDENTIFIERS: + if batch.column(SpecialFields.ROW_ID.name).null_count: + raise ValueError( + "Cannot build global index because _ROW_ID is null.") + for offset in range(0, batch.num_rows, ADD_BATCH_SIZE): + _write_vector_batch( + writer, batch.slice(offset, ADD_BATCH_SIZE), + self._index_columns[0], index_range) + else: + for value, row_id in _extract_index_rows( + batch, + self._index_columns[0], + SpecialFields.ROW_ID.name, + index_range, + ): + writer.write(value, row_id - index_range.from_) + del batch + + if writer is None: + return None + + index_adds = _to_index_manifest_entries( + self._table, + index_split.partition, + index_range, + index_field.id, + self._index_type, + writer.finish(), + ) + if not index_adds: + return None + return CommitMessage( + partition=tuple(index_split.partition.values), + bucket=0, + new_files=[], + index_adds=index_adds, + ) + except BaseException: + if writer is not None: + self._delete_writer_output(writer, index_path) + raise + finally: + if writer is not None: + writer.close() + + def _delete_writer_output(self, writer, index_path: str) -> None: + file_name = getattr(writer, "file_name", None) + if file_name: + self._table.file_io.delete_quietly( + "%s/%s" % (index_path.rstrip("/"), file_name)) + + def _delete_uncommitted_indexes(self, messages) -> None: + path_factory = self._table.path_factory().global_index_path_factory() + for message in messages: + for index_add in message.index_adds: + index_file = index_add.index_file + file_path = ( + index_file.external_path + if index_file.external_path is not None + else path_factory.to_path(index_file.file_name) ) - return messages + self._table.file_io.delete_quietly(file_path) def _create_generic_index_writer(self, index_path: str, index_field): if self._index_type in VINDEX_IDENTIFIERS: diff --git a/paimon-python/pypaimon/tests/global_index_build_test.py b/paimon-python/pypaimon/tests/global_index_build_test.py index 9beac060bb..70985b95fa 100644 --- a/paimon-python/pypaimon/tests/global_index_build_test.py +++ b/paimon-python/pypaimon/tests/global_index_build_test.py @@ -21,6 +21,8 @@ from decimal import Decimal import os import struct import sys +import threading +import time import types import tempfile from unittest.mock import Mock, patch @@ -56,6 +58,7 @@ from pypaimon.tests.data_evolution_test_helpers import ( ) from pypaimon.table.row.generic_row import GenericRow from pypaimon.utils.range import Range +from pypaimon.write.commit_message import CommitMessage class _FakeFile: @@ -1262,6 +1265,7 @@ class GenericIndexStreamingTest(unittest.TestCase): self.builder._table = Mock() self.builder._core_options = Mock() self.builder._core_options.global_index_row_count_per_shard.return_value = 10 + self.builder._core_options.global_index_build_parallelism.return_value = 1 self.builder._index_columns = ['embedding'] self.builder._index_type = 'ivf-flat' self.writer = Mock() @@ -1385,6 +1389,197 @@ class GenericIndexStreamingTest(unittest.TestCase): if failure != 'finish': self.writer.finish.assert_not_called() + def test_build_parallelism_defaults_to_one_and_accepts_override(self): + from pypaimon.common.options.core_options import CoreOptions + from pypaimon.common.options.options import Options + + self.assertEqual( + 1, CoreOptions(Options({})).global_index_build_parallelism()) + self.assertEqual( + 3, + CoreOptions(Options({ + 'global-index.build.parallelism': '3', + })).global_index_build_parallelism(), + ) + + def test_build_parallelism_must_be_positive(self): + self.builder._core_options.global_index_build_parallelism.return_value = 0 + with self.assertRaisesRegex( + ValueError, "global-index.build.parallelism.*greater than 0"): + self._build([self._batch([[10.0]], [10])]) + + def test_parallel_build_is_bounded_and_keeps_shard_order(self): + self.builder._core_options.global_index_build_parallelism.return_value = 2 + shards = [ + (_FakeSplit([]), Range(start, start + 9)) + for start in (0, 10, 20, 30) + ] + lock = threading.Lock() + first_workers = threading.Barrier(2) + active = 0 + max_active = 0 + + def build_shard(_split, row_range, *_args): + nonlocal active, max_active + with lock: + active += 1 + max_active = max(max_active, active) + try: + if row_range.from_ in (0, 10): + first_workers.wait(timeout=2) + time.sleep((40 - row_range.from_) / 1000.0) + return CommitMessage( + partition=(row_range.from_,), + bucket=0, + new_files=[], + ) + finally: + with lock: + active -= 1 + + self.builder._build_generic_shard = Mock(side_effect=build_shard) + module = 'pypaimon.globalindex.create_global_index' + with patch(module + '._split_by_global_index_shard', return_value=shards): + messages = self.builder._build_generic_index( + [], [], Mock(), self.read, '/unused') + + self.assertEqual(2, max_active) + self.assertEqual( + [(0,), (10,), (20,), (30,)], + [message.partition for message in messages], + ) + + def test_parallel_failure_deletes_completed_uncommitted_indexes(self): + self.builder._core_options.global_index_build_parallelism.return_value = 2 + shards = [ + (_FakeSplit([]), Range(0, 9)), + (_FakeSplit([]), Range(10, 19)), + ] + first_started = threading.Event() + failure_released = threading.Event() + index_file = types.SimpleNamespace( + external_path='/index/completed.index', + file_name='completed.index', + ) + completed = CommitMessage( + partition=(), + bucket=0, + new_files=[], + index_adds=[types.SimpleNamespace(index_file=index_file)], + ) + + def build_shard(_split, row_range, *_args): + if row_range.from_ == 0: + first_started.set() + self.assertTrue(failure_released.wait(timeout=2)) + return completed + self.assertTrue(first_started.wait(timeout=2)) + failure_released.set() + raise RuntimeError('injected shard failure') + + self.builder._build_generic_shard = Mock(side_effect=build_shard) + module = 'pypaimon.globalindex.create_global_index' + with patch(module + '._split_by_global_index_shard', return_value=shards), \ + self.assertRaisesRegex(RuntimeError, 'injected shard failure'): + self.builder._build_generic_index( + [], [], Mock(), self.read, '/unused') + + self.builder._table.file_io.delete_quietly.assert_called_once_with( + '/index/completed.index') + + def test_submit_failure_deletes_indexes_built_before_the_failure(self): + # ThreadPoolExecutor.submit() enqueues the work item before starting an + # extra worker, so a submission that raises can still have its shard run + # on an already running worker. Rollback must cover those outputs. + self.builder._core_options.global_index_build_parallelism.return_value = 2 + shards = [ + (_FakeSplit([]), Range(start, start + 9)) + for start in (0, 10, 20, 30) + ] + started = threading.Semaphore(0) + release = threading.Event() + + def build_shard(_split, row_range, *_args): + started.release() + self.assertTrue(release.wait(timeout=5)) + index_file = types.SimpleNamespace( + external_path='/index/shard-%d.index' % row_range.from_, + file_name='shard-%d.index' % row_range.from_, + ) + return CommitMessage( + partition=(row_range.from_,), + bucket=0, + new_files=[], + index_adds=[types.SimpleNamespace(index_file=index_file)], + ) + + self.builder._build_generic_shard = Mock(side_effect=build_shard) + + real_start = threading.Thread.start + submits = [] + + def failing_start(thread): + if thread.name.startswith('paimon-global-index-build'): + submits.append(thread.name) + if len(submits) == 2: + # The first worker already holds a queued shard; releasing it + # here proves rollback sees work accepted before the failure. + self.assertTrue(started.acquire(timeout=5)) + release.set() + raise RuntimeError("can't start new thread") + return real_start(thread) + + module = 'pypaimon.globalindex.create_global_index' + with patch(module + '._split_by_global_index_shard', return_value=shards), \ + patch.object(threading.Thread, 'start', failing_start), \ + self.assertRaisesRegex(RuntimeError, "can't start new thread"): + self.builder._build_generic_index( + [], [], Mock(), self.read, '/unused') + + deleted = { + call[0][0] + for call in self.builder._table.file_io.delete_quietly.call_args_list + } + # Every shard that produced an index file is rolled back, and nothing + # that never ran is deleted. + built = { + '/index/shard-%d.index' % call[0][1].from_ + for call in self.builder._build_generic_shard.call_args_list + } + self.assertTrue(built) + self.assertEqual(built, deleted) + + def test_serial_failure_deletes_completed_uncommitted_indexes(self): + shards = [ + (_FakeSplit([]), Range(0, 9)), + (_FakeSplit([]), Range(10, 19)), + ] + index_file = types.SimpleNamespace( + external_path='/index/completed.index', + file_name='completed.index', + ) + completed = CommitMessage( + partition=(), + bucket=0, + new_files=[], + index_adds=[types.SimpleNamespace(index_file=index_file)], + ) + + def build_shard(_split, row_range, *_args): + if row_range.from_ == 0: + return completed + raise RuntimeError('injected shard failure') + + self.builder._build_generic_shard = Mock(side_effect=build_shard) + module = 'pypaimon.globalindex.create_global_index' + with patch(module + '._split_by_global_index_shard', return_value=shards), \ + self.assertRaisesRegex(RuntimeError, 'injected shard failure'): + self.builder._build_generic_index( + [], [], Mock(), self.read, '/unused') + + self.builder._table.file_io.delete_quietly.assert_called_once_with( + '/index/completed.index') + if __name__ == "__main__": unittest.main() diff --git a/paimon-python/pypaimon/tests/vindex_batch_write_test.py b/paimon-python/pypaimon/tests/vindex_batch_write_test.py index def8858c04..6de7f2ec1d 100644 --- a/paimon-python/pypaimon/tests/vindex_batch_write_test.py +++ b/paimon-python/pypaimon/tests/vindex_batch_write_test.py @@ -169,8 +169,10 @@ class VindexBatchWriteTest(unittest.TestCase): def test_null_row_ids_are_rejected_before_writing_source_batch(self): builder = object.__new__(GlobalIndexBuilder) + builder._table = Mock() builder._core_options = Mock() builder._core_options.global_index_row_count_per_shard.return_value = 10 + builder._core_options.global_index_build_parallelism.return_value = 1 builder._index_type = 'ivf-flat' builder._index_columns = ['embedding'] writer = Mock()
