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 772399a22d [python] Build global index shards concurrently (#9798)
772399a22d is described below
commit 772399a22dc5c81f4abc111cd8d411de838f1031
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()
