[python] Support primary-key vector search with Ray

docs/docs/pypaimon/multimodal-search.md

@@ -199,9 +199,16 @@ data. Historical snapshot and deletion semantics remain the same. Plain
 ## Distributed Vector Search
 
 Use `execution="ray"` to execute vector queries across Ray workers and return
-Arrow results to the driver. This supports single and batch vector queries on
-data-evolution tables; primary-key tables are not supported by this execution
-mode. Full-text and hybrid queries are described below.
+Arrow results to the driver. Single and batch vector queries are supported on
+data-evolution tables and, with the Python read backend, on primary-key tables.
+Batch queries return one Arrow table per input vector, in input order.
+Full-text and hybrid queries are described below.
+
+For primary-key tables, enable the Python read backend with
+`read.native.enabled=false` and `scan.native-plan.enabled=false`. Search
+parallelism is limited by the number of matching buckets; refinement and result
+lookup run on the driver. Primary-key vector queries do not support Native Ray
+execution, Ray Dataset results (`to_ray()`), `with_score()`, or `order_by_score()`.
 
 Install the same PyPaimon and index dependencies on the driver and workers:
 
@@ -261,17 +268,17 @@ batch_neighbors = (
 )
 ```
 
-Each batch task handles all query vectors for one split: index workers reuse
-an open shard across bounded query blocks, and raw workers stream each data
+For data-evolution tables, each batch task handles all query vectors for one split:
+index workers reuse an open shard across bounded query blocks, and raw workers stream each data
 split once with a separate top-k for each query. For refinement, the driver
 selects each query's global candidates, then workers stream their union once per
 read split, scoring each row only for the queries that selected it. Shared final
 row lookup runs on the driver. The entire batch uses one read snapshot.
 Candidate traffic and result memory grow with the number of query vectors;
 split large query collections into smaller batches when necessary.
 
-The driver fixes one read snapshot and plans the query. Workers search individual
-index shards and, when required by the table's search mode, scan unindexed data.
+For data-evolution tables, the driver fixes one read snapshot and plans the query.
+Workers search individual index shards and, when required by the table's search mode, scan unindexed data.
 They return candidate row IDs and scores. For vector queries with
 refinement enabled, the driver first selects the global candidate set, then
 workers read and score those candidate vectors using the index's persisted
@@ -280,8 +287,9 @@ filtering and final row lookup.
 In particular, `pre_filter` filters candidates before ranking; `where()` filters
 the selected rows and can return fewer than the requested number of results.
 
-For indexed vector searches using the Python backend, workers evaluate scalar predicates and build
-live-row filters at the query snapshot. Tasks using the same scalar index files
+For indexed searches on data-evolution tables using the Python backend, workers
+evaluate scalar predicates and build live-row filters at the query snapshot.
+Tasks using the same scalar index files
 share one Ray scalar-index evaluation; its result is passed directly to dependent
 tasks without materializing the bitmap on the driver. Candidate verification
 reads only filter columns and row IDs in the task's shard range, without scanning
@@ -302,15 +310,15 @@ All workers must be able to access the table's storage. Local filesystem paths
 are suitable for a local Ray cluster; multiple nodes require shared storage.
 Each index task searches one shard, while its native index I/O settings still
 apply. Raw-scan parallelism is limited by the number of planned read splits,
-controlled by the table's `source.split.target-size` option. Refinement also
-uses planned read splits; workers return scores for merging.
+controlled by the table's `source.split.target-size` option. For data-evolution
+tables, refinement also uses planned read splits; workers return scores for merging.
 Final row lookup runs on the driver. Candidate traffic grows
 with the number of index shards and the configured refinement budget, so Ray
 execution is most useful when shard search or raw scanning outweighs scheduling
 and transfer costs. Small queries can be faster locally.
 
-Local and Ray index searches merge shard results incrementally in plan order.
-The concurrency limit bounds running tasks plus completed results waiting for
+For data-evolution tables, local and Ray index searches merge shard results
+incrementally in plan order. The concurrency limit bounds running tasks plus completed results waiting for
 earlier shards, so a slow shard can delay further submissions. Global top-k
 selection still uses all merged candidate scores, whose storage grows with the
 number of unique candidate rows per query.

paimon-python/pypaimon/execution/python/read/search/distributed_primary_key_vector_search.py

@@ -0,0 +1,149 @@
+# 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.
+
+"""Bucket task execution with the local PK reader's global ranking semantics."""
+
+from pypaimon.execution.python.read.search.distributed_search import closing_task_results
+from pypaimon.execution.python.read.search.global_index_live_row_filter import table_at_snapshot
+from pypaimon.execution.python.read.search.primary_key_vector_read import PrimaryKeyVectorRead, _top_k
+from pypaimon.execution.python.read.search.primary_key_batch_vector_read import BatchPrimaryKeyVectorRead
+from pypaimon.table.source.search_execution import SearchTask
+from pypaimon.table.source.search_types import PrimaryKeyVectorScanPlan
+
+
+class DistributedPrimaryKeyVectorRead(PrimaryKeyVectorRead):
+
+    def __init__(self, reader, executor):
+        super().__init__(reader._table, reader._limit, reader._vector_column, reader._query_vector,
+                         reader._filter, reader._partition_filter, reader._options)
+        self._executor = executor
+
+    def _context(self, plan):
+        # Workers consume only the supplied bucket metadata, never a fresh scan.
+        return (table_at_snapshot(self._table, plan.snapshot()), self._limit,
+                self._vector_column, self._query_vector, self._options)
+
+    def _search_indexed_candidates(self, plan, indexed_limit):
+        context = self._context(plan)
+        work = [SearchTask(_search_index_bucket, context,
+                           (plan.snapshot_id, split, indexed_limit))
+                for split in plan.splits() if split.payloads]
+        if not work:
+            return []
+
+        def candidates():
+            with closing_task_results(self._executor.map(work)) as tasks:
+                for _, (metric, positions) in tasks:
+                    # Even an empty shard must participate in metric validation.
+                    if metric is not None:
+                        self._set_index_metric(metric)
+                    yield from positions
+
+        # Retain the ANN budget until global refinement, not just each bucket's K.
+        return _top_k(candidates(), indexed_limit)
+
+    def _search_exact_candidates(self, plan):
+        context = self._context(plan)
+        work = [SearchTask(_search_exact_bucket, context,
+                           (plan.snapshot_id, split, self._index_metric))
+                for split in plan.splits() if split.uncovered_data_files]
+        if not work:
+            return []
+
+        def candidates():
+            with closing_task_results(self._executor.map(work)) as tasks:
+                for _, positions in tasks:
+                    yield from positions
+
+        return _top_k(candidates(), self._limit)
+
+
+def _search_index_bucket(context, item):
+    reader = _new_reader(context)
+    snapshot_id, split, indexed_limit = item
+    plan = PrimaryKeyVectorScanPlan(snapshot_id, [split])
+    positions = reader._search_indexed_candidates(plan, indexed_limit)
+    return reader._index_metric, positions
+
+
+def _search_exact_bucket(context, item):
+    reader = _new_reader(context)
+    snapshot_id, split, metric = item
+    if metric is not None:
+        reader._set_index_metric(metric)
+    return reader._search_exact_candidates(PrimaryKeyVectorScanPlan(snapshot_id, [split]))
+
+
+def _new_reader(context):
+    table, limit, column, query, options = context
+    return PrimaryKeyVectorRead(table, limit, column, query, options=options)
+
+
+class DistributedBatchPrimaryKeyVectorRead(BatchPrimaryKeyVectorRead):
+
+    def __init__(self, reader, executor):
+        super().__init__(reader._table, reader._limit, reader._vector_column, reader._query_vectors,
+                         reader._filter, reader._partition_filter, reader._options)
+        self._executor = executor
+
+    def _search_indexed_candidates(self, plan, indexed_limit):
+        return self._search_buckets(plan, _search_batch_index_bucket, indexed_limit, indexed=True)
+
+    def _search_exact_candidates(self, plan):
+        return self._search_buckets(plan, _search_batch_exact_bucket, self._limit, indexed=False)
+
+    def _search_buckets(self, plan, operation, limit, indexed):
+        from itertools import chain
+
+        context = (table_at_snapshot(self._table, plan.snapshot()), self._limit,
+                   self._vector_column, self._query_vectors, self._options)
+        work = [SearchTask(operation, context, (
+            plan.snapshot_id, split, limit if indexed else self._index_metric))
+            for split in plan.splits() if (split.payloads if indexed else split.uncovered_data_files)]
+        candidates = [[] for _ in self._query_vectors]
+        if not work:
+            return candidates
+        with closing_task_results(self._executor.map(work)) as tasks:
+            for _, (metric, results) in tasks:
+                if metric is not None:
+                    self._set_index_metric(metric)
+                if len(results) != len(candidates):
+                    raise ValueError("PK bucket task returned the wrong query count.")
+                for i, positions in enumerate(results):
+                    candidates[i] = _top_k(chain(candidates[i], positions), limit)
+        return candidates
+
+
+def _search_batch_index_bucket(context, item):
+    reader = _new_batch_reader(context)
+    snapshot_id, split, indexed_limit = item
+    positions = reader._search_indexed_candidates(PrimaryKeyVectorScanPlan(snapshot_id, [split]), indexed_limit)
+    return reader._index_metric, positions
+
+
+def _search_batch_exact_bucket(context, item):
+    reader = _new_batch_reader(context)
+    snapshot_id, split, metric = item
+    if metric is not None:
+        reader._set_index_metric(metric)
+    positions = reader._search_exact_candidates(PrimaryKeyVectorScanPlan(snapshot_id, [split]))
+    return reader._index_metric, positions
+
+
+def _new_batch_reader(context):
+    table, limit, column, queries, options = context
+    return BatchPrimaryKeyVectorRead(table, limit, column, queries, options=options)

paimon-python/pypaimon/execution/python/read/search/distributed_search.py

@@ -38,9 +38,6 @@ def closing_task_results(results):
 
 def create_distributed_reader(reader, executor):
     executor = executor.with_concurrency(executor.concurrency)
-    from pypaimon.execution.python.read.search.vector_search_read import (
-        DataEvolutionVectorRead, BatchVectorSearchReadImpl,
-    )
     from pypaimon.execution.python.read.search.full_text_read import DataEvolutionFullTextRead
 
     if isinstance(reader, DataEvolutionFullTextRead):
@@ -49,15 +46,11 @@ def create_distributed_reader(reader, executor):
             raise ValueError("Ray full-text search supports only data-evolution tables.")
         from pypaimon.execution.python.read.search.distributed_full_text_search import DistributedFullTextRead
         return DistributedFullTextRead(reader, executor)
-    if type(reader) not in (DataEvolutionVectorRead, BatchVectorSearchReadImpl):
-        raise ValueError("Ray vector search supports only data-evolution tables.")
-    if not reader._table.options.data_evolution_enabled():
-        raise ValueError("Ray vector search supports only data-evolution tables.")
-    from pypaimon.execution.python.read.search.distributed_vector_search import (
-        DistributedVectorSearchRead, DistributedBatchVectorSearchRead,
-    )
-    if type(reader) is DataEvolutionVectorRead:
-        if not all(math.isfinite(float(value)) for value in reader._query_vector):
-            raise ValueError("Ray vector search requires a finite query vector.")
-        return DistributedVectorSearchRead(reader, executor)
-    return DistributedBatchVectorSearchRead(reader, executor)
+    if getattr(reader, '_new_distributed_reader', None) is None:
+        raise ValueError("This reader does not support distributed search.")
+    return reader._new_distributed_reader(executor)
+
+
+def validate_query_vectors(queries):
+    if not all(math.isfinite(float(value)) for query in queries for value in query):
+        raise ValueError("Ray vector search requires a finite query vector.")

paimon-python/pypaimon/execution/python/read/search/factories.py

@@ -17,46 +17,52 @@
 
 """Construct Python search executors from the public builder configuration."""
 
+from typing import NamedTuple
 
-def new_vector_scan(request):
-    from pypaimon.execution.python.read.search.vector_search_scan import DataEvolutionVectorScan
-    scan_class = DataEvolutionVectorScan
+
+class _VectorSearchTypes(NamedTuple):
+    scan: type
+    read: type
+    batch_read: type
+
+
+def _vector_search_types(request):
     if request._is_primary_key_vector_search():
         from pypaimon.execution.python.read.search.primary_key_vector_scan import PrimaryKeyVectorScan
-        scan_class = PrimaryKeyVectorScan
-    return scan_class(
-        request._table,
-        request._vector_column,
-        filter_=request._filter,
-        partition_filter=request._partition_filter,
-        options=request._options,
+        from pypaimon.execution.python.read.search.primary_key_vector_read import PrimaryKeyVectorRead
+        from pypaimon.execution.python.read.search.primary_key_batch_vector_read import BatchPrimaryKeyVectorRead
+        return _VectorSearchTypes(PrimaryKeyVectorScan, PrimaryKeyVectorRead, BatchPrimaryKeyVectorRead)
+    from pypaimon.execution.python.read.search.vector_search_scan import DataEvolutionVectorScan
+    from pypaimon.execution.python.read.search.vector_search_read import (
+        DataEvolutionVectorRead, BatchVectorSearchReadImpl,
     )
+    return _VectorSearchTypes(DataEvolutionVectorScan, DataEvolutionVectorRead, BatchVectorSearchReadImpl)
 
 
-def new_vector_read(request):
-    from pypaimon.execution.python.read.search.vector_search_read import DataEvolutionVectorRead
-    read_class = DataEvolutionVectorRead
-    if request._is_primary_key_vector_search():
-        from pypaimon.execution.python.read.search.primary_key_vector_read import PrimaryKeyVectorRead
-        read_class = PrimaryKeyVectorRead
-    return read_class(
+def new_vector_scan(request):
+    return _vector_search_types(request).scan(
         request._table,
-        request._limit,
         request._vector_column,
-        request._query_vector,
         filter_=request._filter,
         partition_filter=request._partition_filter,
         options=request._options,
     )
 
 
+def new_vector_read(request):
+    return _new_vector_reader(request, _vector_search_types(request).read, request._query_vector)
+
+
 def new_batch_vector_read(request):
-    from pypaimon.execution.python.read.search.vector_search_read import BatchVectorSearchReadImpl
-    return BatchVectorSearchReadImpl(
+    return _new_vector_reader(request, _vector_search_types(request).batch_read, request._query_vectors)
+
+
+def _new_vector_reader(request, read_class, queries):
+    return read_class(
         request._table,
         request._limit,
         request._vector_column,
-        request._query_vectors,
+        queries,
         filter_=request._filter,
         partition_filter=request._partition_filter,
         options=request._options,