diff --git a/docs/docs/pypaimon/multimodal-search.md b/docs/docs/pypaimon/multimodal-search.md index 1c56a5e04bff..0601aded02b7 100644 --- a/docs/docs/pypaimon/multimodal-search.md +++ b/docs/docs/pypaimon/multimodal-search.md @@ -50,6 +50,32 @@ docs.create_index("embedding", index_type="ivf-pq") docs.create_index("content", index_type="full-text") ``` +### Build vector indexes with Ray + +Pass `execution="ray"` to distribute native vector-index shards across Ray +workers. BTree, Bitmap and full-text builds still use local execution. + +```python +docs.create_index( + "embedding", index_type="ivf-flat", + options={"global-index.row-count-per-shard": "1000000"}, + execution="ray", concurrency=4, ray_remote_args={"num_cpus": 1}, +) +``` + +The driver plans uncovered row ranges from one snapshot. Workers build files; +the driver publishes their index entries in one commit after all shards succeed. +`concurrency` bounds running shards; `global-index.build.parallelism` controls +only local thread-based builds. Install the same `pypaimon[ray,vindex]` +dependencies on every worker and use storage accessible to the entire cluster. + +Write tasks require `max_retries=0` and `retry_exceptions=False` (set by default). +On a task failure, the driver stops submissions, waits for dispatched tasks, +and removes their uncommitted output files. Rerun the build after the failure; +already committed index coverage is skipped. Driver/cluster termination can +still leave orphan files for normal orphan-file maintenance. Retain the source +snapshot's files until the build and its commit finish. + ## Search Use `search` for one vector query or one full-text query. diff --git a/paimon-python/pypaimon/filesystem/jindo_file_system_handler.py b/paimon-python/pypaimon/filesystem/jindo_file_system_handler.py index 467ba03927ba..5a242df8dec3 100644 --- a/paimon-python/pypaimon/filesystem/jindo_file_system_handler.py +++ b/paimon-python/pypaimon/filesystem/jindo_file_system_handler.py @@ -235,6 +235,14 @@ def __init__(self, root_path: str, catalog_options: Options): config = build_jindo_config(catalog_options) self._jindo_fs = jfs.connect(self.root_path, "root", config) + def __getstate__(self): + # Native clients and thread locks belong to this process. Transport only + # connection configuration; immutable-file size hints can be refilled. + return {"root_path": self.root_path, "properties": self.properties} + + def __setstate__(self, state): + self.__init__(state["root_path"], state["properties"]) + def register_file_size(self, path: str, file_size: int): """Register an immutable file's size supplied by Paimon metadata.""" normalized = self._normalize_path(path) diff --git a/paimon-python/pypaimon/globalindex/create_global_index.py b/paimon-python/pypaimon/globalindex/create_global_index.py index 59261011480c..569f2f8861f4 100644 --- a/paimon-python/pypaimon/globalindex/create_global_index.py +++ b/paimon-python/pypaimon/globalindex/create_global_index.py @@ -73,6 +73,10 @@ def create_global_index( partition_filter: Optional[Predicate] = None, partitions: Optional[Union[Dict[str, object], Sequence[Dict[str, object]]]] = None, options: Optional[Dict[str, object]] = None, + *, + execution="local", + concurrency=None, + ray_remote_args=None, ) -> int: """Build and commit global index files for a table. @@ -87,7 +91,8 @@ def create_global_index( partitions=partitions, options=options, ) - messages = builder.build() + messages = builder.build( + execution=execution, concurrency=concurrency, ray_remote_args=ray_remote_args) if not messages: return 0 @@ -159,7 +164,15 @@ def __init__( if self._index_type in _GENERIC_INDEX_IDENTIFIERS: self._validate_generic_index_table() - def build(self) -> List[CommitMessage]: + def build(self, *, execution="local", concurrency=None, ray_remote_args=None) -> List[CommitMessage]: + if execution not in ("local", "ray"): + raise ValueError("execution must be 'local' or 'ray'.") + if execution == "local" and (concurrency is not None or ray_remote_args is not None): + raise ValueError("Ray options require execution='ray'.") + if execution == "ray": + from pypaimon.ray.vector_index_build import validate_build_options + concurrency, ray_remote_args = validate_build_options( + self, concurrency, ray_remote_args) read_builder = self._table.new_read_builder() partition_filter = self._resolve_partition_filter() if partition_filter is not None: @@ -203,6 +216,11 @@ def build(self) -> List[CommitMessage]: if self._index_type in _SORTED_INDEX_IDENTIFIERS: return self._build_sorted_index( splits, unindexed_ranges, index_field, table_read, index_path) + if execution == "ray": + from pypaimon.ray.vector_index_build import build_vector_index + return build_vector_index( + self, splits, unindexed_ranges, index_field, table_read, + index_path, concurrency, ray_remote_args) return self._build_generic_index( splits, unindexed_ranges, index_field, table_read, index_path) @@ -380,7 +398,7 @@ def build_shard(index_split, index_range): raise def _build_generic_shard( - self, index_split, index_range, index_field, table_read, index_path: str + self, index_split, index_range, index_field, table_read, index_path: str, file_name=None ) -> Optional[CommitMessage]: from pypaimon.read.table_read import _ClosableArrowBatchReader @@ -396,6 +414,8 @@ def _build_generic_shard( if writer is None: writer = self._create_generic_index_writer( index_path, index_field) + if file_name is not None: + writer.file_name = file_name if self._index_type in VINDEX_IDENTIFIERS: if batch.column(SpecialFields.ROW_ID.name).null_count: raise ValueError( diff --git a/paimon-python/pypaimon/multimodal/table.py b/paimon-python/pypaimon/multimodal/table.py index 6c0ff50ff87f..00f1c03c4e60 100644 --- a/paimon-python/pypaimon/multimodal/table.py +++ b/paimon-python/pypaimon/multimodal/table.py @@ -534,9 +534,11 @@ def search_hybrid( pre_filter=pre_filter, ) - def create_index(self, column, index_type, options=None): + def create_index(self, column, index_type, options=None, *, execution="local", + concurrency=None, ray_remote_args=None): return self.raw_table.create_global_index( - column, index_type=_normalize_index_type(index_type), options=options) + column, index_type=_normalize_index_type(index_type), options=options, + execution=execution, concurrency=concurrency, ray_remote_args=ray_remote_args) def _merge( self, diff --git a/paimon-python/pypaimon/ray/vector_index_build.py b/paimon-python/pypaimon/ray/vector_index_build.py new file mode 100644 index 000000000000..15575aed1975 --- /dev/null +++ b/paimon-python/pypaimon/ray/vector_index_build.py @@ -0,0 +1,81 @@ +# 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. + +"""Ray shard construction with one driver-side global-index commit.""" + +from pypaimon.globalindex.build_plan import split_by_global_index_shard +from pypaimon.globalindex.index_file_utils import new_global_index_file_name +from pypaimon.globalindex.vindex.vindex_vector_global_index_reader import VINDEX_IDENTIFIERS +from pypaimon.ray.vector_search import _execution_options + + +def validate_build_options(builder, concurrency, ray_remote_args): + if builder._index_type not in VINDEX_IDENTIFIERS: + raise ValueError("Ray index building supports only native vector indexes.") + concurrency, options = _execution_options(concurrency, ray_remote_args) + # Retrying a write task can race with a lost worker writing the same file. + # A failed build can instead be rerun after its uncommitted outputs close. + if options.get("max_retries", 0) != 0 or options.get("retry_exceptions", False): + raise ValueError("Ray index building requires max_retries=0 and retry_exceptions=False.") + options.update(max_retries=0, retry_exceptions=False) + if builder._core_options.global_index_row_count_per_shard() <= 0: + raise ValueError("Option 'global-index.row-count-per-shard' must be greater than 0.") + return concurrency, options + + +def build_vector_index(builder, splits, unindexed_ranges, index_field, table_read, + index_path, concurrency, remote_args): + import ray + + shards = split_by_global_index_shard( + splits, builder._core_options.global_index_row_count_per_shard(), unindexed_ranges) + # Allocate names before dispatch, so failed/lost result delivery cannot hide + # files from cleanup. Workers never publish index manifests or snapshots. + names = [new_global_index_file_name("vindex-" + builder._index_type) for _ in shards] + context = ray.put((builder, index_field, table_read, index_path)) + remote = ray.remote(_build_shard).options(**remote_args) + pending = {} + results = {} + next_shard = 0 + try: + while pending or next_shard < len(shards): + while len(pending) < concurrency and next_shard < len(shards): + ref = remote.remote(context, shards[next_shard], names[next_shard]) + pending[ref] = next_shard + next_shard += 1 + ready, _ = ray.wait(list(pending), num_returns=1) + ref = ready[0] + results[pending.pop(ref)] = ray.get(ref) + return [results[i] for i in range(len(shards)) if results[i] is not None] + except BaseException: + # Do not delete a file while another worker may still be writing it. + # Stop submissions and drain every dispatched task before cleanup. + for ref in pending: + try: + ray.get(ref) + except BaseException: + pass + for name in names: + builder._table.file_io.delete_quietly(index_path.rstrip("/") + "/" + name) + raise + + +def _build_shard(context, shard, file_name): + builder, index_field, table_read, index_path = context + index_split, index_range = shard + return builder._build_generic_shard( + index_split, index_range, index_field, table_read, index_path, file_name=file_name) diff --git a/paimon-python/pypaimon/table/file_store_table.py b/paimon-python/pypaimon/table/file_store_table.py index 76f2d019a1d8..8dea054608a1 100644 --- a/paimon-python/pypaimon/table/file_store_table.py +++ b/paimon-python/pypaimon/table/file_store_table.py @@ -461,7 +461,8 @@ def new_batch_vector_search_builder(self) -> 'BatchVectorSearchBuilder': def create_global_index(self, index_column, index_type: str = "btree", partition_filter=None, partitions=None, - options: Optional[dict] = None) -> int: + options: Optional[dict] = None, *, execution="local", + concurrency=None, ray_remote_args=None) -> int: from pypaimon.globalindex.create_global_index import \ create_global_index return create_global_index( @@ -471,6 +472,7 @@ def create_global_index(self, index_column, index_type: str = "btree", partition_filter=partition_filter, partitions=partitions, options=options, + execution=execution, concurrency=concurrency, ray_remote_args=ray_remote_args, ) def drop_global_index(self, index_column, index_type: str = "btree", diff --git a/paimon-python/pypaimon/tests/jindo_file_system_test.py b/paimon-python/pypaimon/tests/jindo_file_system_test.py index a3d7bd1194e6..0c95c9228bed 100644 --- a/paimon-python/pypaimon/tests/jindo_file_system_test.py +++ b/paimon-python/pypaimon/tests/jindo_file_system_test.py @@ -17,6 +17,8 @@ import io import os +import pickle +import threading import types import unittest import uuid @@ -102,6 +104,40 @@ def tell(self): stream = JindoInputFile(JindoStream()) self.assertEqual(5, stream.seek(0, io.SEEK_END)) + def test_pickle_reconnects_without_native_client_lock_or_cache(self): + options = Options({ + "fs.oss.impl": "jindo", "fs.oss.endpoint": "https://oss-cn-hangzhou.aliyuncs.com", + "fs.oss.accessKeyId": "test-ak", "fs.oss.accessKeySecret": "test-sk", + "fs.oss.securityToken": "test-token", "fs.jindocache.client.metrics.enable": True, + }) + # A native client is not serializable even after the handler lock is + # removed. Give this stand-in its own lock to enforce reconstruction. + clients = [types.SimpleNamespace(lock=threading.Lock()) for _ in range(2)] + connect = mock.Mock(side_effect=clients) + with mock.patch.object(jindo_module, "JINDO_AVAILABLE", True), \ + mock.patch.object(jindo_module, "jfs", types.SimpleNamespace(connect=connect)), \ + mock.patch.object(jindo_module, "jutil", types.SimpleNamespace(Config=_RecordingConfig)): + original = JindoFileSystemHandler("oss://bucket/", options) + original.register_file_size("object", 123) + # Exercise PyArrow's enclosing filesystem serialization as well. + encoded = pickle.dumps(PyFileSystem(original)) + with original._known_file_sizes_lock: + restored = pickle.loads(encoded).handler + self.assertIs(original._jindo_fs, clients[0]) + self.assertIs(restored._jindo_fs, clients[1]) + self.assertEqual(restored.root_path, original.root_path) + self.assertEqual(restored.properties.to_map(), options.to_map()) + self.assertIsNot(restored._known_file_sizes_lock, original._known_file_sizes_lock) + self.assertIsNone(restored._known_file_size("oss://bucket/object")) + restored.register_file_size("object", 456) + self.assertEqual(restored._known_file_size("oss://bucket/object"), 456) + self.assertEqual(original._known_file_size("oss://bucket/object"), 123) + self.assertEqual(connect.call_count, 2) + for call in connect.call_args_list: + self.assertEqual(call[0][:2], ("oss://bucket/", "root")) + self.assertEqual(call[0][2].values["fs.oss.securityToken"], "test-token") + self.assertEqual(call[0][2].values["fs.jindocache.client.metrics.enable"], "true") + def test_forwards_native_options_to_connect(self): created_config = _RecordingConfig() config_factory = mock.Mock(return_value=created_config) diff --git a/paimon-python/pypaimon/tests/multimodal_table_test.py b/paimon-python/pypaimon/tests/multimodal_table_test.py index f202fb2629d9..cd60ffc97f07 100644 --- a/paimon-python/pypaimon/tests/multimodal_table_test.py +++ b/paimon-python/pypaimon/tests/multimodal_table_test.py @@ -2835,7 +2835,7 @@ def test_create_index_normalizes_full_text_alias(self): calls = [] - def create_global_index(column, index_type, options=None): + def create_global_index(column, index_type, options=None, **kwargs): calls.append((column, index_type, options)) return index_type @@ -2865,7 +2865,7 @@ def test_create_index_requires_index_type(self): options=_PARQUET_OPTIONS, ) - def create_global_index(column, index_type, options=None): + def create_global_index(column, index_type, options=None, **kwargs): return index_type docs.raw_table.create_global_index = create_global_index diff --git a/paimon-python/pypaimon/tests/ray_vector_index_build_test.py b/paimon-python/pypaimon/tests/ray_vector_index_build_test.py new file mode 100644 index 000000000000..af5f6bb62184 --- /dev/null +++ b/paimon-python/pypaimon/tests/ray_vector_index_build_test.py @@ -0,0 +1,191 @@ +# 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. + +"""Exercise Ray vector-index construction with real worker processes.""" + +import os +import threading +import time +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +ray = pytest.importorskip("ray") +pytest.importorskip("paimon_vindex") + +from pypaimon.globalindex.create_global_index import GlobalIndexBuilder +from pypaimon.index.index_file_handler import IndexFileHandler +from pypaimon.ray import vector_index_build +from pypaimon.tests import ray_vector_search_test as fixtures +from pypaimon.tests.ray_vector_search_test import add_rows, VECTORS + +ray_cluster = fixtures.ray_cluster +table = fixtures.table + + +OPTIONS = {"global-index.row-count-per-shard": "3", "ivf-flat.nlist": "1"} + + +def index_files(table): + return [entry.index_file for entry in IndexFileHandler(table.raw_table).scan( + table.raw_table.snapshot_manager().get_latest_snapshot()) + if entry.index_file.index_type == "ivf-flat"] + + +@pytest.mark.parametrize("concurrency", [1, 2]) +def test_build_and_incremental_coverage(table, ray_cluster, concurrency): + add_rows(table, VECTORS) + table.delete("id = 0") + before = table.raw_table.snapshot_manager().get_latest_snapshot().id + assert table.create_index("embedding", "ivf-flat", options=OPTIONS, + execution="ray", concurrency=concurrency) == 2 + assert table.raw_table.snapshot_manager().get_latest_snapshot().id == before + 1 + files = sorted(index_files(table), key=lambda f: f.global_index_meta.row_range_start) + assert [(f.global_index_meta.row_range_start, f.global_index_meta.row_range_end) for f in files] == [(0, 2), (3, 5)] + query = table.search([1., 1.], column="embedding", options={"ivf.nprobe": "1"}).select(["id"]).limit(2) + assert query.to_list() == query.to_list(execution="ray") + assert all(row["id"] != 0 for row in query.to_list()) + assert table.create_index("embedding", "ivf-flat", options=OPTIONS, execution="ray") == 0 + add_rows(table, [[20., 1.]], 6) + assert table.create_index("embedding", "ivf-flat", options=OPTIONS, execution="ray") == 1 + assert len(index_files(table)) == 3 + + +def test_build_messages_do_not_publish_or_include_later_appends(table, ray_cluster): + add_rows(table, VECTORS[:3]) + builder = GlobalIndexBuilder(table.raw_table, "embedding", "ivf-flat", options=OPTIONS) + snapshot = table.raw_table.snapshot_manager().get_latest_snapshot().id + messages = builder.build(execution="ray") + assert table.raw_table.snapshot_manager().get_latest_snapshot().id == snapshot + assert not index_files(table) + add_rows(table, VECTORS[3:], 3) + commit = table.raw_table.new_batch_write_builder().new_commit() + try: + commit.commit(messages) + finally: + commit.close() + assert [(f.global_index_meta.row_range_start, f.global_index_meta.row_range_end) + for f in index_files(table)] == [(0, 2)] + assert table.create_index("embedding", "ivf-flat", options=OPTIONS, execution="ray") == 1 + + +def _write_then_fail(context, shard, name): + if shard[1].from_: + time.sleep(0.3) + message = vector_index_build._build_shard(context, shard, name) + if not shard[1].from_: + raise RuntimeError("failed after writing an index file") + return message + + +def test_worker_failure_drains_and_cleans_all_outputs(table, ray_cluster): + add_rows(table, VECTORS) + before = table.raw_table.snapshot_manager().get_latest_snapshot().id + # Pass a real serialized worker function; a failure after finish() must also + # clean files whose commit messages were never delivered to the driver. + original_remote = ray.remote + + def remote(worker): + return original_remote(_write_then_fail) + + with patch.object(ray, "remote", remote): + with pytest.raises(Exception, match="failed after writing"): + table.create_index("embedding", "ivf-flat", options=OPTIONS, + execution="ray", concurrency=2) + assert table.raw_table.snapshot_manager().get_latest_snapshot().id == before + assert not index_files(table) + root = table.raw_table.path_factory().global_index_path_factory().global_index_root_path() + from pathlib import Path + assert not list(Path(root).glob("*.index")) + assert table.create_index("embedding", "ivf-flat", options=OPTIONS, execution="ray") == 2 + + +@pytest.mark.parametrize("kwargs,match", [ + ({"execution": "other"}, "execution"), + ({"concurrency": 1}, "execution='ray'"), + ({"execution": "ray", "concurrency": 0}, "concurrency"), + ({"execution": "ray", "ray_remote_args": {"max_retries": 1}}, "max_retries"), + ({"execution": "ray", "ray_remote_args": {"retry_exceptions": True}}, "max_retries"), +]) +def test_validate_before_empty_plan(table, kwargs, match): + with pytest.raises(ValueError, match=match): + table.create_index("embedding", "ivf-flat", options=OPTIONS, **kwargs) + + +def test_reject_non_vector_ray_build(table): + with pytest.raises(ValueError, match="only native vector"): + table.create_index("id", "btree", execution="ray") + + +@contextmanager +def _jindo_sdk_stub(): + from pypaimon.filesystem import jindo_file_system_handler as jindo_module + from pypaimon.filesystem import pyarrow_file_io as arrow_module + from pypaimon.tests.jindo_file_system_test import _RecordingConfig + + def connect(root, user, config): + return SimpleNamespace(pid=os.getpid(), lock=threading.Lock(), config=config.values) + + with patch.object(jindo_module, "JINDO_AVAILABLE", True), \ + patch.object(arrow_module, "JINDO_AVAILABLE", True), \ + patch.object(jindo_module, "jfs", SimpleNamespace(connect=connect)), \ + patch.object(jindo_module, "jutil", SimpleNamespace(Config=_RecordingConfig)): + yield + + +def _inspect_jindo_context(context_refs): + # Nested object refs postpone deserialization until the worker has installed + # the SDK stand-in. Only Config/connect are replaced, not serialization. + with _jindo_sdk_stub(): + builder, field, table_read, index_path = ray.get(context_refs[0]) + handler = builder._table.file_io.filesystem.handler + assert table_read.table.file_io is builder._table.file_io + assert field.name == "embedding" + assert handler._known_file_size("oss://bucket/object") is None + handler.register_file_size("object", 456) + return (os.getpid(), handler._jindo_fs.pid, handler._known_file_size("oss://bucket/object"), + handler.root_path, handler._jindo_fs.config["fs.oss.securityToken"], index_path) + + +def test_jindo_build_context_reconstructs_on_worker(table, ray_cluster): + from pypaimon.common.options import Options + from pypaimon.filesystem.oss_file_io import OssFileIO + from pypaimon.read.table_read import TableRead + from pypaimon.table.special_fields import SpecialFields + + builder = GlobalIndexBuilder(table.raw_table, "embedding", "ivf-flat", options=OPTIONS) + field = table.raw_table.field_dict["embedding"] + reader = TableRead(table.raw_table, None, [field, SpecialFields.ROW_ID]) + with _jindo_sdk_stub(): + file_io = OssFileIO("oss://bucket/table", Options({ + "fs.oss.impl": "jindo", "fs.oss.endpoint": "oss-cn-hangzhou.aliyuncs.com", + "fs.oss.accessKeyId": "test-ak", "fs.oss.accessKeySecret": "test-sk", + "fs.oss.securityToken": "test-token", + })) + original = file_io.filesystem.handler + original.register_file_size("object", 123) + with patch.object(table.raw_table, "file_io", file_io): + # This is the exact tuple broadcast by build_vector_index. + context = ray.put((builder, field, reader, "oss://bucket/index")) + result = ray.get(ray.remote(_inspect_jindo_context).remote([context])) + worker_pid, client_pid, size, root, token, index_path = result + assert worker_pid != os.getpid() + assert client_pid == worker_pid + assert (size, root, token, index_path) == (456, "oss://bucket/", "test-token", "oss://bucket/index") + assert original._known_file_size("oss://bucket/object") == 123