From da87cce08afa4602c2cc998f50d49373292dfb39 Mon Sep 17 00:00:00 2001 From: chaoyang Date: Mon, 14 Sep 2026 11:41:43 +0800 Subject: [PATCH 1/2] [python] Build global index shards concurrently --- .../global-index/manage-indexes.mdx | 1 + .../pypaimon/common/options/core_options.py | 14 ++ .../globalindex/create_global_index.py | 205 +++++++++++++----- .../pypaimon/tests/global_index_build_test.py | 133 ++++++++++++ .../pypaimon/tests/vindex_batch_write_test.py | 2 + 5 files changed, 299 insertions(+), 56 deletions(-) diff --git a/docs/docs/multimodal-table/global-index/manage-indexes.mdx b/docs/docs/multimodal-table/global-index/manage-indexes.mdx index 2e1ff0c9f9cf..24f2e70a5484 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 903697a94ff2..a94effa11ec3 100644 --- a/paimon-python/pypaimon/common/options/core_options.py +++ b/paimon-python/pypaimon/common/options/core_options.py @@ -903,6 +903,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() @@ -1620,6 +1631,9 @@ def global_index_thread_num(self) -> Optional[int]: 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 d68221a0ae6c..0d2a086068d7 100644 --- a/paimon-python/pypaimon/globalindex/create_global_index.py +++ b/paimon-python/pypaimon/globalindex/create_global_index.py @@ -17,6 +17,7 @@ """Build global index files from Python.""" +from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait from functools import cmp_to_key from typing import Dict, List, Optional, Sequence, Union @@ -288,72 +289,164 @@ def _create_sorted_index_writer(self, index_path: str, key_serializer): 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: + 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 + + futures = [] + try: + with ThreadPoolExecutor( + max_workers=min(parallelism, len(shards)), + thread_name_prefix="paimon-global-index-build", + ) as executor: + futures = [ + executor.submit( + self._build_generic_shard, + index_split, + index_range, + index_field, + table_read, + index_path, + ) + for index_split, index_range in shards + ] + 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, + ) + 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. Delete every completed index because build() + # will not return commit messages after a shard failure. + messages = [] + for future in futures: + if future.cancelled() or not future.done(): + continue + try: + message = future.result() + except BaseException: continue + if message is not None: + messages.append(message) + self._delete_uncommitted_indexes(messages) + 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 - index_adds = _to_index_manifest_entries( - self._table, - index_split.partition, - index_range, - index_field.id, - self._index_type, - writer.finish(), - ) - 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, - ) + 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 9beac060bb2d..b72ab8ec756e 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 @@ 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.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 @@ def setUp(self): 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,135 @@ def batches(): 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_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 def8858c0475..6de7f2ec1dc8 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 @@ def test_builder_filters_ranges_before_vector_validation(self): 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() From e40b2bde3e573add17d01b88716d4bc0324b61da Mon Sep 17 00:00:00 2001 From: chaoyang Date: Tue, 15 Sep 2026 11:39:32 +0800 Subject: [PATCH 2/2] [python] Roll back global index shards accepted before a submit failure The parallel build assigned `futures` only after the whole comprehension succeeded, so a submission that raised (for example `RuntimeError: can't start new thread`) left it as the empty list. Shards accepted before the failure still ran during executor shutdown, but rollback saw no futures and left their index files behind. Appending each future as it is returned is not enough: `submit()` puts the work item on the queue before it starts an extra worker, so even the submission that raises can run its shard on an already running worker. Track output ownership in the worker itself so rollback no longer depends on the future list, and cancel still-queued shards when submission fails. Co-Authored-By: Claude Opus 5 --- .../globalindex/create_global_index.py | 55 +++++++++------- .../pypaimon/tests/global_index_build_test.py | 62 +++++++++++++++++++ 2 files changed, 93 insertions(+), 24 deletions(-) diff --git a/paimon-python/pypaimon/globalindex/create_global_index.py b/paimon-python/pypaimon/globalindex/create_global_index.py index 0d2a086068d7..93473bb7b3e6 100644 --- a/paimon-python/pypaimon/globalindex/create_global_index.py +++ b/paimon-python/pypaimon/globalindex/create_global_index.py @@ -17,6 +17,7 @@ """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 @@ -319,23 +320,38 @@ def _build_generic_index( 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: - futures = [ - executor.submit( - self._build_generic_shard, - index_split, - index_range, - index_field, - table_read, - index_path, - ) - for index_split, index_range in shards - ] + 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 @@ -353,19 +369,10 @@ def _build_generic_index( 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. Delete every completed index because build() - # will not return commit messages after a shard failure. - messages = [] - for future in futures: - if future.cancelled() or not future.done(): - continue - try: - message = future.result() - except BaseException: - continue - if message is not None: - messages.append(message) - self._delete_uncommitted_indexes(messages) + # 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( diff --git a/paimon-python/pypaimon/tests/global_index_build_test.py b/paimon-python/pypaimon/tests/global_index_build_test.py index b72ab8ec756e..70985b95fa18 100644 --- a/paimon-python/pypaimon/tests/global_index_build_test.py +++ b/paimon-python/pypaimon/tests/global_index_build_test.py @@ -1487,6 +1487,68 @@ def build_shard(_split, row_range, *_args): 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)),