[python] Build global index shards concurrently

Author: TheR1sing3un

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. |

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)
 

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: