[python] Return search results through distributed Ray row lookup

docs/docs/pypaimon/multimodal-search.md

@@ -378,3 +378,30 @@ batch_neighbors = (
     .to_list()
 )
 ```
+
+## Return search rows as a Ray Dataset
+
+Vector queries support `to_ray()` to run candidate search on Ray and fetch the
+selected columns on workers. Single queries return one Dataset; `search_vectors`
+returns one Dataset per input vector. Candidate IDs and scores are merged on the
+driver, while projected row data stays in Ray blocks.
+
+```python
+ds = (docs.search([0.1, 0.2, 0.3], column="embedding")
+      .select(["content", "image"])
+      .with_score().order_by_score().limit(100)
+      .to_ray(concurrency=4, override_num_blocks=4))
+```
+
+Search runs eagerly; row lookup is lazy and uses the same captured snapshot.
+Retain that snapshot's files until downstream Dataset actions finish. `where`
+filters the candidates during lookup without refilling top-k. `order_by_score`
+performs a distributed global sort, breaking ties by ascending row ID. Empty
+results retain their output schema, including the score column.
+
+BLOB values are serialized descriptors, as with `scan().to_ray()`. Pass the
+Dataset to `docs.map_with_blobs(...)` to resolve payloads on workers. Projections
+must be nonempty and have unique names. This API requires data evolution;
+`execution="local"` selects local candidate search while keeping distributed
+lookup. Batch queries share candidate-search work but have separate lazy lookup
+pipelines for each query.

paimon-python/pypaimon/multimodal/query.py

@@ -586,8 +586,29 @@ def read_blobs(self, *args, **kwargs):
     def stream_blobs(self, *args, **kwargs):
         raise TypeError("stream_blobs is only supported on scan(), not search queries.")
 
-    def to_ray(self, *args, **kwargs):
-        raise TypeError("to_ray is only supported on scan(), not search queries.")
+    def to_ray(self, *, execution="ray", concurrency=None, ray_remote_args=None,
+               override_num_blocks=None):
+        """Return search rows as a Ray Dataset, with BLOB descriptors.
+
+        Search resolves candidates eagerly on one snapshot. Projected row data
+        is fetched lazily by Ray workers on that same snapshot. Retain its files
+        until the Dataset has finished executing.
+        """
+        from pypaimon.ray.search_result import read_search_result
+        from pypaimon.ray.vector_search import _execution_options
+
+        concurrency, remote_args = _execution_options(concurrency, ray_remote_args)
+        query = self._for_execution()
+        if execution == "ray":
+            result = query._execute_ray(concurrency, remote_args)
+        elif execution == "local":
+            result = query._result_factory(query)
+        else:
+            raise ValueError("execution must be 'local' or 'ray'.")
+        return read_search_result(query, result, concurrency, remote_args, override_num_blocks)
+
+    def _execute_ray(self, concurrency, ray_remote_args):
+        raise NotImplementedError("Ray search is not supported for this query type.")
 
     def to_arrow_batch_reader(self, *args, **kwargs):
         raise TypeError(
@@ -652,6 +673,12 @@ def to_list(self, *, execution="local", concurrency=None, ray_remote_args=None)
             execution=execution, concurrency=concurrency,
             ray_remote_args=ray_remote_args).to_pylist()
 
+    def _execute_ray(self, concurrency, ray_remote_args):
+        from pypaimon.ray.vector_search import _execute_vector_search
+        return _execute_vector_search(
+            self._vector_search_builder(self), concurrency=concurrency,
+            ray_remote_args=ray_remote_args)
+
     def _execute_vector(self, query):
         return self._vector_search_builder(query).execute_local()
 
@@ -779,6 +806,26 @@ def to_arrow(self, *, execution="local", concurrency=None, ray_remote_args=None)
             concurrency=concurrency, ray_remote_args=ray_remote_args)
         return query._read_batch_results(results)
 
+    def to_ray(self, *, execution="ray", concurrency=None, ray_remote_args=None,
+               override_num_blocks=None):
+        """Return one lazy Ray Dataset per query, with distributed row lookup."""
+        from pypaimon.ray.batch_vector_search import _execute_batch_vector_search
+        from pypaimon.ray.search_result import read_search_result
+        from pypaimon.ray.vector_search import _execution_options
+
+        concurrency, remote_args = _execution_options(concurrency, ray_remote_args)
+        query = self._for_execution()
+        if execution == "ray":
+            results = _execute_batch_vector_search(
+                query._batch_vector_search_builder(query), concurrency=concurrency,
+                ray_remote_args=remote_args)
+        elif execution == "local":
+            results = query._execute_batch_vector(query)
+        else:
+            raise ValueError("execution must be 'local' or 'ray'.")
+        return [read_search_result(query, result, concurrency, remote_args, override_num_blocks)
+                for result in results]
+
     def _read_batch_results(self, results):
         from pypaimon.globalindex.global_index_result import GlobalIndexResult
         from pypaimon.utils.roaring_bitmap import RoaringBitmap64

paimon-python/pypaimon/ray/search_result.py

@@ -0,0 +1,91 @@
+# 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.
+
+"""Snapshot-pinned distributed lookup of global-index search candidates."""
+
+from copy import copy
+import uuid
+
+import pyarrow as pa
+
+from pypaimon.table.special_fields import SpecialFields
+
+
+def read_search_result(query, result, concurrency, remote_args, override_num_blocks):
+    import ray
+
+    query._require_row_id_scores()
+    lookup = copy(query)
+    lookup._table = query._table.copy({"blob-as-descriptor": "true"})
+    projection = query._effective_projection()
+    lookup._projection = (list(projection) if projection is not None
+                          else [field.name for field in query._table.fields])
+    if not lookup._projection or len(set(lookup._projection)) != len(lookup._projection):
+        raise ValueError("Ray search requires a nonempty projection with unique column names.")
+    row_id = SpecialFields.ROW_ID.name
+    added_row_id = row_id not in lookup._projection
+    if added_row_id:
+        lookup._projection.append(row_id)
+    lookup._limit = None  # Candidates already contain the global top-k.
+    builder = lookup._configured_read_builder()
+    reader = builder.new_read()
+    if query._metadata_only_result():
+        dataset = ray.data.from_arrow(pa.table({row_id: pa.array(list(result.results()), type=pa.int64())}))
+    else:
+        splits = builder.new_scan().with_global_index_result(result).plan().splits()
+        dataset = reader.to_ray(
+            splits, concurrency=concurrency, ray_remote_args=remote_args,
+            override_num_blocks=override_num_blocks)
+
+    # Only compact candidate metadata crosses the driver. Payloads remain in
+    # worker blocks, including during score attachment and global ordering.
+    rank_column = "_paimon_search_rank_" + uuid.uuid4().hex
+    finish = copy(query)
+    finish._sort_by_score = False
+    sort_by_score = query._sort_by_score
+    score_getter = (result.score_getter()
+                    if sort_by_score and not result.results().is_empty() else None)
+
+    def finish_batch(batch):
+        output = finish._finish_search_result(batch, result, False)
+        if sort_by_score:
+            output = output.append_column(rank_column, pa.array(
+                [-score_getter(value) for value in batch[row_id].to_pylist()], type=pa.float64()))
+        return output
+
+    empty = reader._output_arrow_schema().empty_table()
+    dataset = dataset.map_batches(finish_batch, batch_format="pyarrow")
+    empty = finish_batch(empty)
+    if query._sort_by_score:
+        dataset = dataset.sort([rank_column, row_id]).drop_columns([rank_column])
+        empty = empty.drop_columns([rank_column])
+    if added_row_id:
+        dataset = dataset.drop_columns([row_id])
+        empty = empty.drop_columns([row_id])
+    # Sort/map may discard every empty block. Restore the final typed schema.
+    dataset = dataset.union(ray.data.from_arrow(empty))
+    if sort_by_score:
+        # Projection tasks after Sort can finish out of order. Preserve the
+        # sorted block order through the whole returned Dataset, without
+        # changing the process-wide DataContext or unrelated Datasets.
+        dataset.context.execution_options.preserve_order = True
+    setattr(dataset, "_paimon_blob_file_io", lookup._table.file_io)
+    setattr(dataset, "_paimon_blob_columns", query._readable_blob_columns())
+    maps, arrays = query._nested_blob_columns()
+    setattr(dataset, "_paimon_map_blob_columns", maps)
+    setattr(dataset, "_paimon_array_blob_columns", arrays)
+    return dataset

paimon-python/pypaimon/tests/multimodal_table_test.py

@@ -1922,8 +1922,6 @@ def test_search_query_rejects_blob_reads(self):
             t.search([1.0, 0.0, 0.0], column="emb").read_blobs("img")
         with self.assertRaisesRegex(TypeError, "only supported on scan"):
             t.search([1.0, 0.0, 0.0], column="emb").stream_blobs("img")
-        with self.assertRaisesRegex(TypeError, "only supported on scan"):
-            t.search([1.0, 0.0, 0.0], column="emb").to_ray()
 
     def test_scan_stream_blobs(self):
         obs = self.conn.create_table(

paimon-python/pypaimon/tests/ray_search_order_test.py

@@ -0,0 +1,72 @@
+# 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.
+
+import time
+from unittest.mock import patch
+
+import pytest
+
+ray = pytest.importorskip("ray")
+
+from ray.data.dataset import Dataset
+from pypaimon.tests import ray_vector_search_test as fixtures
+
+table = fixtures.table
+
+
+@pytest.fixture(scope="module")
+def ray_cluster():
+    started = not ray.is_initialized()
+    if started:
+        ray.init(address="local", num_cpus=4, include_dashboard=False,
+                 object_store_memory=1024 * 1024 * 1024)
+    yield
+    if started:
+        ray.shutdown()
+
+
+def _slow_first_drop(self, columns, **kwargs):
+    def drop(batch):
+        if batch.num_rows and batch["id"][0].as_py() == 0:
+            time.sleep(1)
+        return batch.drop_columns(columns)
+    return self.map_batches(drop, batch_format="pyarrow", zero_copy_batch=True, **kwargs)
+
+
+@pytest.mark.parametrize("batch", [False, True])
+def test_score_order_survives_slow_projection(table, ray_cluster, batch):
+    # Four physical files and score ties exercise global order and the row-id
+    # tie breaker. Delay the first sorted partition, leaving its data unchanged.
+    for start in range(0, 24, 6):
+        fixtures.add_rows(table, [[float(i // 2), 0.] for i in range(start, start + 6)], start)
+    global_context = ray.data.DataContext.get_current()
+    original_order = global_context.execution_options.preserve_order
+    unrelated = ray.data.range(1)
+    unrelated_order = unrelated.context.execution_options.preserve_order
+    query = table.search_vectors([[0., 0.]]) if batch else table.search([0., 0.])
+    query.select(["id"]).with_score().order_by_score().limit(24)
+    with patch.object(Dataset, "drop_columns", _slow_first_drop):
+        result = query.to_ray(execution="local", override_num_blocks=4)
+    dataset = result[0] if batch else result
+    # Prevent tiny test blocks being coalesced into a single projection task.
+    dataset.context.target_min_block_size = 1
+    assert [row["id"] for row in dataset.take_all()] == list(range(24))
+    assert dataset.take(1)[0]["id"] == 0
+    assert global_context.execution_options.preserve_order == original_order
+    assert unrelated.context.execution_options.preserve_order == unrelated_order
+    unordered = table.search([0., 0.]).select(["id"]).to_ray(execution="local")
+    assert unordered.context.execution_options.preserve_order == original_order

paimon-python/pypaimon/tests/ray_search_result_test.py

@@ -0,0 +1,99 @@
+# 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.
+
+from unittest.mock import patch
+
+import pytest
+
+ray = pytest.importorskip("ray")
+
+from pypaimon.tests import ray_vector_search_test as fixtures
+from pypaimon.multimodal.query import _PreFilterQuery
+
+ray_cluster = fixtures.ray_cluster
+table = fixtures.table
+
+
+@pytest.mark.parametrize("execution", ["local", "ray"])
+@pytest.mark.parametrize("scores", [False, True])
+def test_distributed_lookup_and_snapshot(table, ray_cluster, execution, scores):
+    fixtures.add_rows(table, fixtures.VECTORS[:3])
+    fixtures.add_rows(table, fixtures.VECTORS[3:], 3)
+    query = table.search([1., 1.]).select(["id", "category"]).limit(4)
+    if scores:
+        query.with_score().order_by_score()
+    expected = query.to_list()
+    with patch.object(_PreFilterQuery, "_read_global_index_result", side_effect=AssertionError("driver lookup")):
+        dataset = query.to_ray(execution=execution, concurrency=2, override_num_blocks=2)
+    table.delete("id = 0")
+    table.update("id = 1", {"category": "changed"})
+    rows = dataset.take_all()
+    if scores:
+        assert rows == expected
+    else:
+        assert sorted(rows, key=lambda r: r["id"]) == sorted(expected, key=lambda r: r["id"])
+    assert dataset.schema().names == (["id", "category", "_score"] if scores else ["id", "category"])
+
+
+def test_filters_empty_scores_and_metadata_only(table, ray_cluster):
+    fixtures.add_rows(table, fixtures.VECTORS)
+    query = table.search([1., 1.], pre_filter="category = 'yes'").select(["id"]).with_score().order_by_score()
+    assert query.to_ray().take_all() == query.to_list()
+    query.where("id = -1")
+    dataset = query.to_ray()
+    assert dataset.take_all() == []
+    assert dataset.schema().names == ["id", "_score"]
+    query = table.search([1., 1.]).select(["_ROW_ID"]).with_score().order_by_score().limit(3)
+    assert query.to_ray().take_all() == query.to_list()
+
+
+def test_batch_search_datasets(table, ray_cluster):
+    fixtures.add_rows(table, fixtures.VECTORS)
+    query = table.search_vectors([[1., 1.], [4., 1.]]).select(["id"]).with_score().order_by_score().limit(2)
+    assert [ds.take_all() for ds in query.to_ray()] == query.to_list()
+
+
+def test_blob_descriptors_can_be_resolved_on_workers(tmp_path, ray_cluster):
+    import pyarrow as pa
+    import pypaimon.multimodal as pm
+
+    table = pm.connect(options={"warehouse": str(tmp_path)}).create_table(
+        "images", schema=pa.schema([("id", pa.int64()), ("image", pa.large_binary()),
+                                    ("embedding", pa.list_(pa.float32(), 2))]),
+        options={"file.format": "parquet", "vector.file.format": "parquet",
+                 "vector-index.search-mode": "full"})
+    table.add([{"id": 0, "image": b"a", "embedding": [1., 1.]},
+               {"id": 1, "image": b"b", "embedding": [4., 1.]}])
+    dataset = table.search([1., 1.]).select(["id", "image"]).limit(1).to_ray()
+
+    def resolve(scalar, blobs):
+        return scalar.append_column("body", pa.array(blobs["image"]))
+
+    assert table.map_with_blobs(dataset, ["image"], resolve).take_all() == [{"id": 0, "body": b"a"}]
+
+
+@pytest.mark.parametrize("projection", [["id"], ["_ROW_ID"]])
+@pytest.mark.parametrize("execution", ["local", "ray"])
+def test_empty_table_retains_scored_schema(table, ray_cluster, projection, execution):
+    query = table.search([1., 1.]).select(projection).with_score().order_by_score()
+    dataset = query.to_ray(execution=execution)
+    assert dataset.take_all() == []
+    assert dataset.schema().names == projection + ["_score"]
+    batch = table.search_vectors([[1., 1.]]).select(projection).with_score().order_by_score()
+    dataset = batch.to_ray(execution=execution)[0]
+    assert dataset.take_all() == []
+    assert dataset.schema().names == projection + ["_score"]