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 787f16cbce [python] Return search results through distributed Ray row
lookup (#10429)
787f16cbce is described below
commit 787f16cbce5ca74fa6cddb050a0af06d98e04a0a
Author: chaoyang <[email protected]>
AuthorDate: Fri Oct 9 21:58:24 2026 +0800
[python] Return search results through distributed Ray row lookup (#10429)
---
docs/docs/pypaimon/multimodal-search.md | 27 ++++++
paimon-python/pypaimon/multimodal/query.py | 51 ++++++++++-
paimon-python/pypaimon/ray/search_result.py | 91 ++++++++++++++++++++
.../pypaimon/tests/multimodal_table_test.py | 2 -
.../pypaimon/tests/ray_search_order_test.py | 72 ++++++++++++++++
.../pypaimon/tests/ray_search_result_test.py | 99 ++++++++++++++++++++++
6 files changed, 338 insertions(+), 4 deletions(-)
diff --git a/docs/docs/pypaimon/multimodal-search.md
b/docs/docs/pypaimon/multimodal-search.md
index b317e75180..c63dc59a99 100644
--- a/docs/docs/pypaimon/multimodal-search.md
+++ b/docs/docs/pypaimon/multimodal-search.md
@@ -410,3 +410,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.
diff --git a/paimon-python/pypaimon/multimodal/query.py
b/paimon-python/pypaimon/multimodal/query.py
index d4d70267c8..195c05e9ac 100644
--- a/paimon-python/pypaimon/multimodal/query.py
+++ b/paimon-python/pypaimon/multimodal/query.py
@@ -586,8 +586,29 @@ class _PreFilterQuery(ScanQuery):
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 @@ class VectorQuery(_PreFilterQuery):
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()
@@ -817,6 +844,26 @@ class BatchVectorQuery(_PreFilterQuery):
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
diff --git a/paimon-python/pypaimon/ray/search_result.py
b/paimon-python/pypaimon/ray/search_result.py
new file mode 100644
index 0000000000..12fe70501b
--- /dev/null
+++ b/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
diff --git a/paimon-python/pypaimon/tests/multimodal_table_test.py
b/paimon-python/pypaimon/tests/multimodal_table_test.py
index f202fb2629..7f5028ed47 100644
--- a/paimon-python/pypaimon/tests/multimodal_table_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_table_test.py
@@ -1922,8 +1922,6 @@ class MultimodalTableTest(unittest.TestCase):
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(
diff --git a/paimon-python/pypaimon/tests/ray_search_order_test.py
b/paimon-python/pypaimon/tests/ray_search_order_test.py
new file mode 100644
index 0000000000..858b5750d2
--- /dev/null
+++ b/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
+
+
[email protected](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)
+
+
[email protected]("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
diff --git a/paimon-python/pypaimon/tests/ray_search_result_test.py
b/paimon-python/pypaimon/tests/ray_search_result_test.py
new file mode 100644
index 0000000000..01b53e067c
--- /dev/null
+++ b/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
+
+
[email protected]("execution", ["local", "ray"])
[email protected]("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"}]
+
+
[email protected]("projection", [["id"], ["_ROW_ID"]])
[email protected]("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"]
