diff --git a/docs/docs/pypaimon/multimodal-search.md b/docs/docs/pypaimon/multimodal-search.md index 1c56a5e04bff..854580c7bed1 100644 --- a/docs/docs/pypaimon/multimodal-search.md +++ b/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. diff --git a/paimon-python/pypaimon/multimodal/query.py b/paimon-python/pypaimon/multimodal/query.py index e743bdf6ae7f..4349ca8c55ea 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/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 diff --git a/paimon-python/pypaimon/ray/search_result.py b/paimon-python/pypaimon/ray/search_result.py new file mode 100644 index 000000000000..12fe70501b4a --- /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 f202fb2629d9..7f5028ed4712 100644 --- a/paimon-python/pypaimon/tests/multimodal_table_test.py +++ b/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( 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 000000000000..858b5750d21c --- /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 + + +@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 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 000000000000..01b53e067c8c --- /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 + + +@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"]