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 6f13b1bbb3 [python] Support distributed vector index construction with
Ray (#10426)
6f13b1bbb3 is described below
commit 6f13b1bbb3c7640800adcfba7b8549a8735b7602
Author: chaoyang <[email protected]>
AuthorDate: Fri Oct 9 21:58:54 2026 +0800
[python] Support distributed vector index construction with Ray (#10426)
---
docs/docs/pypaimon/multimodal-search.md | 26 +++
.../filesystem/jindo_file_system_handler.py | 8 +
.../pypaimon/globalindex/create_global_index.py | 26 ++-
paimon-python/pypaimon/multimodal/table.py | 6 +-
paimon-python/pypaimon/ray/vector_index_build.py | 81 +++++++++
paimon-python/pypaimon/table/file_store_table.py | 4 +-
.../pypaimon/tests/jindo_file_system_test.py | 36 ++++
.../pypaimon/tests/multimodal_table_test.py | 4 +-
.../pypaimon/tests/ray_vector_index_build_test.py | 191 +++++++++++++++++++++
9 files changed, 374 insertions(+), 8 deletions(-)
diff --git a/docs/docs/pypaimon/multimodal-search.md
b/docs/docs/pypaimon/multimodal-search.md
index c63dc59a99..4470bf3d4e 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 467ba03927..5a242df8de 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 @@ class JindoFileSystemHandler(FileSystemHandler):
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 5926101148..569f2f8861 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 @@ class GlobalIndexBuilder:
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 @@ class GlobalIndexBuilder:
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 @@ class GlobalIndexBuilder:
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 @@ class GlobalIndexBuilder:
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 6c0ff50ff8..00f1c03c4e 100644
--- a/paimon-python/pypaimon/multimodal/table.py
+++ b/paimon-python/pypaimon/multimodal/table.py
@@ -534,9 +534,11 @@ class MultimodalTable:
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 0000000000..15575aed19
--- /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 76f2d019a1..8dea054608 100644
--- a/paimon-python/pypaimon/table/file_store_table.py
+++ b/paimon-python/pypaimon/table/file_store_table.py
@@ -461,7 +461,8 @@ class FileStoreTable(Table):
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 @@ class FileStoreTable(Table):
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 a3d7bd1194..0c95c9228b 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 @@ class JindoConfigTest(unittest.TestCase):
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 7f5028ed47..00bd51e7d2 100644
--- a/paimon-python/pypaimon/tests/multimodal_table_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_table_test.py
@@ -2833,7 +2833,7 @@ class MultimodalTableTest(unittest.TestCase):
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
@@ -2863,7 +2863,7 @@ class MultimodalTableTest(unittest.TestCase):
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 0000000000..af5f6bb621
--- /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"]
+
+
[email protected]("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
+
+
[email protected]("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
