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 61b2d64c9d [python] Support Ray options in vector result conversions
(#9998)
61b2d64c9d is described below
commit 61b2d64c9df312d2dc02cfe79a1baf213375109a
Author: chaoyang <[email protected]>
AuthorDate: Sun Sep 20 16:11:15 2026 +0800
[python] Support Ray options in vector result conversions (#9998)
---
docs/docs/pypaimon/multimodal-search.md | 6 ++
paimon-python/pypaimon/multimodal/query.py | 24 +++++-
.../tests/ray_vector_result_conversion_test.py | 87 ++++++++++++++++++++++
3 files changed, 113 insertions(+), 4 deletions(-)
diff --git a/docs/docs/pypaimon/multimodal-search.md
b/docs/docs/pypaimon/multimodal-search.md
index 2fc052bd7f..8c98e2e331 100644
--- a/docs/docs/pypaimon/multimodal-search.md
+++ b/docs/docs/pypaimon/multimodal-search.md
@@ -147,6 +147,12 @@ neighbors = (
the number of in-flight Ray tasks, with a default of 4. `ray_remote_args`
supplies
Ray task options, including resources and retry settings; `num_returns` is
managed by PyPaimon. These two arguments require `execution="ray"`.
+
+For single and batch vector queries, `to_pandas()` and `to_list()` accept the
+same `execution`, `concurrency`, and `ray_remote_args` arguments as
`to_arrow()`.
+For example, use `query.to_pandas(execution="ray", concurrency=4)` to return a
+DataFrame, or one DataFrame per query for a batch. Result conversion runs on
the
+driver after the search and lookup complete.
Single-vector Ray queries require finite query values and fail if stored
vectors produce NaN scores.
diff --git a/paimon-python/pypaimon/multimodal/query.py
b/paimon-python/pypaimon/multimodal/query.py
index c3aa7b1bd2..af2a983f9a 100644
--- a/paimon-python/pypaimon/multimodal/query.py
+++ b/paimon-python/pypaimon/multimodal/query.py
@@ -554,6 +554,18 @@ class VectorQuery(_PreFilterQuery):
concurrency=concurrency, ray_remote_args=ray_remote_args)
return query._read_global_index_result(result)
+ def to_pandas(self, *, execution="local", concurrency=None,
ray_remote_args=None):
+ """Execute with the same options as to_arrow and return a DataFrame."""
+ return self.to_arrow(
+ execution=execution, concurrency=concurrency,
+ ray_remote_args=ray_remote_args).to_pandas()
+
+ def to_list(self, *, execution="local", concurrency=None,
ray_remote_args=None) -> List[dict]:
+ """Execute with the same options as to_arrow and return rows."""
+ return self.to_arrow(
+ execution=execution, concurrency=concurrency,
+ ray_remote_args=ray_remote_args).to_pylist()
+
def _execute_vector(self, query):
return self._vector_search_builder(query).execute_local()
@@ -716,11 +728,15 @@ class BatchVectorQuery(_PreFilterQuery):
output.append(table.take(pa.array(selected, type=pa.int64())))
return output
- def to_pandas(self):
- return [table.to_pandas() for table in self.to_arrow()]
+ def to_pandas(self, *, execution="local", concurrency=None,
ray_remote_args=None):
+ """Return one DataFrame per query, using the same options as
to_arrow."""
+ return [table.to_pandas() for table in self.to_arrow(
+ execution=execution, concurrency=concurrency,
ray_remote_args=ray_remote_args)]
- def to_list(self) -> List[List[dict]]:
- return [table.to_pylist() for table in self.to_arrow()]
+ def to_list(self, *, execution="local", concurrency=None,
ray_remote_args=None) -> List[List[dict]]:
+ """Return rows per query, using the same options as to_arrow."""
+ return [table.to_pylist() for table in self.to_arrow(
+ execution=execution, concurrency=concurrency,
ray_remote_args=ray_remote_args)]
def _execute_batch_vector(self, query):
return self._batch_vector_search_builder(query).execute_batch_local()
diff --git a/paimon-python/pypaimon/tests/ray_vector_result_conversion_test.py
b/paimon-python/pypaimon/tests/ray_vector_result_conversion_test.py
new file mode 100644
index 0000000000..5c6925c7b5
--- /dev/null
+++ b/paimon-python/pypaimon/tests/ray_vector_result_conversion_test.py
@@ -0,0 +1,87 @@
+# 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
+
+pytest.importorskip("ray")
+
+from pypaimon.ray import batch_vector_search, vector_search
+from pypaimon.tests import ray_vector_search_test as fixtures
+from pypaimon.tests.ray_vector_search_test import VECTORS, add_rows,
build_index
+
+ray_cluster = fixtures.ray_cluster
+table = fixtures.table
+
+
+def make_query(table, batch):
+ queries = [[1., 1.], [-1., 2.], [1., 1.]]
+ return (table.search_vectors(queries) if batch else
table.search(queries[0])).select(["id"]).limit(2)
+
+
[email protected]("batch", [False, True])
[email protected]("method", ["to_pandas", "to_list"])
[email protected]("indexed", [False, True])
+def test_ray_result_conversion(table, ray_cluster, batch, method, indexed):
+ add_rows(table, VECTORS)
+ if indexed:
+ build_index(table)
+ add_rows(table, [[1., 1.]], 6)
+ query = make_query(table, batch)
+ expected = query.to_list()
+ module = batch_vector_search if batch else vector_search
+ name = "_execute_batch_vector_search" if batch else
"_execute_vector_search"
+ with patch.object(module, name, wraps=getattr(module, name)) as execute:
+ actual = getattr(query, method)(
+ execution="ray", concurrency=2, ray_remote_args={"num_cpus": 1})
+ execute.assert_called_once()
+ assert execute.call_args[1] == {"concurrency": 2, "ray_remote_args":
{"num_cpus": 1}}
+ if method == "to_pandas":
+ actual = [frame.to_dict("records") for frame in actual] if batch else
actual.to_dict("records")
+ assert actual == expected
+
+
[email protected]("batch", [False, True])
[email protected]("method", ["to_pandas", "to_list"])
+def test_conversion_keeps_empty_results_and_default_local_execution(table,
batch, method):
+ query = make_query(table, batch)
+ with patch.object(vector_search, "_execute_vector_search",
side_effect=AssertionError("Ray dispatch")), \
+ patch.object(batch_vector_search, "_execute_batch_vector_search",
+ side_effect=AssertionError("Ray dispatch")):
+ actual = getattr(query, method)()
+ values = actual if batch else [actual]
+ assert len(values) == (3 if batch else 1)
+ assert all(len(value) == 0 for value in values)
+ if method == "to_pandas":
+ assert all(list(frame.columns) == ["id"] for frame in values)
+
+
[email protected]("batch", [False, True])
[email protected]("method", ["to_pandas", "to_list"])
+def test_conversion_validates_options_and_propagates_failure(table, batch,
method):
+ query = make_query(table, batch)
+ with pytest.raises(ValueError, match="execution must"):
+ getattr(query, method)(execution="invalid")
+ with pytest.raises(ValueError, match="require execution"):
+ getattr(query, method)(concurrency=2)
+ module = batch_vector_search if batch else vector_search
+ name = "_execute_batch_vector_search" if batch else
"_execute_vector_search"
+ failure = RuntimeError("worker failed")
+ with patch.object(module, name, side_effect=failure),
pytest.raises(RuntimeError) as exc:
+ getattr(query, method)(execution="ray")
+ assert exc.value is failure
