From 8218194d1b27dacca6cc0e1a3603da7cee08914b Mon Sep 17 00:00:00 2001 From: tzphh <742429651@qq.com> Date: Tue, 27 Oct 2026 14:15:17 +0800 Subject: [PATCH] [python] Support primary-key vector search with Ray --- docs/docs/pypaimon/multimodal-search.md | 34 +- .../distributed_primary_key_vector_search.py | 149 ++++ .../python/read/search/distributed_search.py | 23 +- .../execution/python/read/search/factories.py | 50 +- .../search/primary_key_batch_vector_read.py | 157 +++++ .../read/search/primary_key_vector_read.py | 112 +-- .../python/read/search/vector_search_read.py | 20 + paimon-python/pypaimon/multimodal/query.py | 12 +- .../distributed_reader_selection_test.py | 178 +++++ .../tests/pk_search_capabilities_test.py | 80 +++ .../ray_primary_key_vector_search_test.py | 645 ++++++++++++++++++ .../tests/ray_search_execution_test.py | 23 +- .../vector_search_factory_selection_test.py | 135 ++++ 13 files changed, 1511 insertions(+), 107 deletions(-) create mode 100644 paimon-python/pypaimon/execution/python/read/search/distributed_primary_key_vector_search.py create mode 100644 paimon-python/pypaimon/execution/python/read/search/primary_key_batch_vector_read.py create mode 100644 paimon-python/pypaimon/tests/distributed_reader_selection_test.py create mode 100644 paimon-python/pypaimon/tests/pk_search_capabilities_test.py create mode 100644 paimon-python/pypaimon/tests/ray_primary_key_vector_search_test.py create mode 100644 paimon-python/pypaimon/tests/vector_search_factory_selection_test.py diff --git a/docs/docs/pypaimon/multimodal-search.md b/docs/docs/pypaimon/multimodal-search.md index 57513fae0f79..713479edbe8d 100644 --- a/docs/docs/pypaimon/multimodal-search.md +++ b/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,8 +268,8 @@ 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 @@ -270,8 +277,8 @@ 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. diff --git a/paimon-python/pypaimon/execution/python/read/search/distributed_primary_key_vector_search.py b/paimon-python/pypaimon/execution/python/read/search/distributed_primary_key_vector_search.py new file mode 100644 index 000000000000..f64703224a34 --- /dev/null +++ b/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) diff --git a/paimon-python/pypaimon/execution/python/read/search/distributed_search.py b/paimon-python/pypaimon/execution/python/read/search/distributed_search.py index 783798d3584a..09878696c218 100644 --- a/paimon-python/pypaimon/execution/python/read/search/distributed_search.py +++ b/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.") diff --git a/paimon-python/pypaimon/execution/python/read/search/factories.py b/paimon-python/pypaimon/execution/python/read/search/factories.py index c9beb96ff0d9..6f81e53a8c1f 100644 --- a/paimon-python/pypaimon/execution/python/read/search/factories.py +++ b/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, diff --git a/paimon-python/pypaimon/execution/python/read/search/primary_key_batch_vector_read.py b/paimon-python/pypaimon/execution/python/read/search/primary_key_batch_vector_read.py new file mode 100644 index 000000000000..3b3e4ae072e6 --- /dev/null +++ b/paimon-python/pypaimon/execution/python/read/search/primary_key_batch_vector_read.py @@ -0,0 +1,157 @@ +# 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. + +"""Batch PK searches sharing index readers and physical vector reads.""" + +import math +from itertools import chain + +import pyarrow as pa + +from pypaimon.common.options.core_options import GlobalIndexSearchMode +from pypaimon.execution.python.read.search.primary_key_vector_read import ( + PrimaryKeyVectorRead, _deleted_positions, _include_row_ids, _localized_candidates, _top_k, +) +from pypaimon.execution.python.read.search.vector_search_read import ( + _iter_arrow_batch_scores, _iter_arrow_scores, +) +from pypaimon.globalindex.batch_vector_search import BatchVectorSearch +from pypaimon.index.pk.primary_key_index_source_meta import PrimaryKeyIndexSourceMeta +from pypaimon.table.source.primary_key_scored_result import ( + PrimaryKeyScoredResult, PrimaryKeySearchPosition, _partition_bytes, _position_key, +) +from pypaimon.table.source.search_types import ( + BatchVectorSearchRead, PrimaryKeyVectorScanPlan, PrimaryKeyVectorSearchSplit, +) + + +class BatchPrimaryKeyVectorRead(PrimaryKeyVectorRead, BatchVectorSearchRead): + + def __init__(self, table, limit, vector_column, query_vectors, filter_=None, + partition_filter=None, options=None): + queries = [list(query) for query in query_vectors] + if not queries or not queries[0] or any(len(query) != len(queries[0]) for query in queries): + raise ValueError("PK batch search requires nonempty query vectors of the same dimension.") + if not all(math.isfinite(float(value)) for query in queries for value in query): + raise ValueError("PK batch search requires finite query vectors.") + super().__init__(table, limit, vector_column, queries[0], filter_, partition_filter, options) + self._query_vectors = queries + + def _new_distributed_reader(self, executor): + if type(self) is not BatchPrimaryKeyVectorRead: + raise ValueError("Custom vector readers must define their own distributed execution.") + from pypaimon.execution.python.read.search.distributed_search import validate_query_vectors + from pypaimon.execution.python.read.search.distributed_primary_key_vector_search import ( + DistributedBatchPrimaryKeyVectorRead, + ) + validate_query_vectors(self._query_vectors) + return DistributedBatchPrimaryKeyVectorRead(self, executor) + + def read_batch_plan(self, plan, executor=None): + if executor is not None: + return self._with_executor(executor).read_batch_plan(plan) + if not isinstance(plan, PrimaryKeyVectorScanPlan): + raise ValueError("Primary-key vector read requires a PrimaryKeyVectorScanPlan.") + self._index_metric = None + index_type = self._table.options.primary_key_vector_index_type(self._vector_column.name) + indexed_limit = self._indexed_search_limit(index_type) + if indexed_limit > 2147483647: + raise ValueError("Vector search limit overflow: limit=%d" % self._limit) + indexed = self._search_indexed_candidates(plan, indexed_limit) + if self._configured_refine_factor(index_type) > 0: + indexed = self._rerank_indexed_batch(plan, indexed, index_type) + exact = ([[] for _ in self._query_vectors] + if self._table.options.vector_index_search_mode() == GlobalIndexSearchMode.FAST + else self._search_exact_candidates(plan)) + sources = [split.data_split for split in plan.splits()] + return [PrimaryKeyScoredResult(plan.snapshot_id, sources, _top_k(chain(ann, raw), self._limit)) + for ann, raw in zip(indexed, exact)] + + def _read_batch(self, splits, snapshot): + if splits and snapshot is None: + raise ValueError("PK batch search requires a captured plan; use read_batch_plan().") + return self.read_batch_plan(PrimaryKeyVectorScanPlan( + 0 if snapshot is None else snapshot.id, splits, snapshot)) + + def _search_indexed_candidates(self, plan, indexed_limit): + candidates = [[] for _ in self._query_vectors] + deleted = _deleted_positions(self._table, plan) + # Keep each payload open across bounded query blocks. + block_size = max(1, min(64, (1 << 20) // len(self._query_vectors[0]))) + for split in plan.splits(): + for payload in split.payloads: + source = PrimaryKeyIndexSourceMeta.from_index_file(payload) + include = _include_row_ids(split, source.source_files, deleted) + count = sum(file.row_count for file in source.source_files) + reader, offset = self._open_offset_reader([payload], 0, count - 1) + try: + for start in range(0, len(self._query_vectors), block_size): + queries = self._query_vectors[start:start + block_size] + search = BatchVectorSearch(queries, indexed_limit, self._vector_column.name, + include_row_ids=include, options=self._options) + results = offset.visit_batch_vector_search(search).result() + if len(results) != len(queries): + raise ValueError("PK index task returned the wrong query count.") + for i, result in enumerate(results, start): + candidates[i] = _top_k(chain(candidates[i], _localized_candidates( + result, split, source, deleted)), indexed_limit) + finally: + reader.close() + return candidates + + def _search_exact_candidates(self, plan): + candidates = [[] for _ in self._query_vectors] + metric = self._search_metric(self._table.options.primary_key_vector_index_type(self._vector_column.name)) + for split, data_file, positions, vectors in self._raw_batches(plan): + partition = _partition_bytes(split.data_split.partition) + for start, query_index, scores in _iter_arrow_batch_scores(vectors, self._query_vectors, metric): + rows = (PrimaryKeySearchPosition( + partition, split.data_split.bucket, data_file.file_name, position, score) + for position, score in zip(positions[start:start + len(scores)], scores) if score is not None) + candidates[query_index] = _top_k(chain(candidates[query_index], rows), self._limit) + return candidates + + def _rerank_indexed_batch(self, plan, candidates, index_type): + selected = [{_position_key(position): position for position in group} for group in candidates] + if any(len(group) != len(positions) for group, positions in zip(selected, candidates)): + raise ValueError("Duplicate primary-key vector candidate.") + union = {key: position for group in selected for key, position in group.items()} + sources = [split.data_split for split in plan.splits()] + result = PrimaryKeyScoredResult(plan.snapshot_id, sources, list(union.values())) + # Reuse the physical reader on the candidate union; each query scores only + # its own globally selected ANN candidates, never another query's rows. + splits = [PrimaryKeyVectorSearchSplit( + split.data_split(), (), (split.files[0].file_name,), + {split.files[0].file_name: tuple(split.row_ranges())}) for split in result.splits] + candidate_plan = PrimaryKeyVectorScanPlan(plan.snapshot_id, splits, plan.snapshot()) + metric = self._search_metric(index_type) + reranked = [[] for _ in self._query_vectors] + for split, data_file, positions, vectors in self._raw_batches(candidate_plan): + keys = [(_partition_bytes(split.data_split.partition), split.data_split.bucket, + data_file.file_name, position) for position in positions] + for i, query in enumerate(self._query_vectors): + offsets = [offset for offset, key in enumerate(keys) if key in selected[i]] + if not offsets: + continue + rows = vectors.take(pa.array(offsets, type=pa.int64())) + for offset, score in zip(offsets, _iter_arrow_scores(rows, query, metric)): + if score is None: + raise ValueError("Primary-key vector candidate contains a null vector.") + reranked[i].append(selected[i].pop(keys[offset]).with_score(score)) + if any(selected): + raise ValueError("Failed to read all primary-key vector candidates for reranking.") + return [_top_k(group, self._limit) for group in reranked] diff --git a/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_read.py b/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_read.py index 227af32445c0..3c7dc898fd5d 100644 --- a/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_read.py +++ b/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_read.py @@ -17,6 +17,7 @@ from concurrent.futures import wait from heapq import nsmallest +from itertools import islice from pypaimon.common.options.core_options import GlobalIndexSearchMode from pypaimon.index.pk.primary_key_index_source_meta import PrimaryKeyIndexSourceMeta @@ -35,6 +36,16 @@ class PrimaryKeyVectorRead(DataEvolutionVectorRead): """Search PK vector payloads and localize payload row ids to data files.""" + def _new_distributed_reader(self, executor): + if type(self) is not PrimaryKeyVectorRead: + raise ValueError("Custom vector readers must define their own distributed execution.") + from pypaimon.execution.python.read.search.distributed_search import validate_query_vectors + from pypaimon.execution.python.read.search.distributed_primary_key_vector_search import ( + DistributedPrimaryKeyVectorRead, + ) + validate_query_vectors([self._query_vector]) + return DistributedPrimaryKeyVectorRead(self, executor) + def read_plan(self, plan, executor=None): if executor is not None: return self._with_executor(executor).read_plan(plan) @@ -48,6 +59,27 @@ def read_plan(self, plan, executor=None): raise ValueError( "Vector search limit overflow: limit=%d, refine factor=%d" % (self._limit, indexed_limit // self._limit)) + # Match Java's separately bounded indexed and exact queues. The ANN + # queue must retain indexed_limit until exact reranking is complete. + indexed_candidates = self._search_indexed_candidates(plan, indexed_limit) + if self._configured_refine_factor(index_type) > 0 and indexed_candidates: + indexed_candidates = self._rerank_indexed( + plan, indexed_candidates, index_type) + else: + indexed_candidates = _top_k(indexed_candidates, self._limit) + exact_candidates = [] + if (self._table.options.vector_index_search_mode() + != GlobalIndexSearchMode.FAST): + exact_candidates = self._search_exact_candidates(plan) + candidates = _top_k( + (candidate + for group in (indexed_candidates, exact_candidates) + for candidate in group), + self._limit) + source_splits = [split.data_split for split in plan.splits()] + return PrimaryKeyScoredResult(plan.snapshot_id, source_splits, candidates) + + def _search_indexed_candidates(self, plan, indexed_limit): futures = [] contexts = [] deleted_positions = _deleted_positions(self._table, plan) @@ -66,42 +98,13 @@ def read_plan(self, plan, executor=None): def indexed_candidate_iter(): for future, context in zip(futures, contexts): - result = future.result() - if result is None: - continue - score_getter = result.score_getter() - for local_row_id in result.results(): - source, row_position = _localize( - context[1].source_files, local_row_id) - if row_position in deleted_positions.get( - (id(context[0]), source.file_name), set()): - continue - if not _allowed(context[0], source.file_name, row_position): - continue - yield PrimaryKeySearchPosition( - _partition_bytes(context[0].data_split.partition), - context[0].data_split.bucket, source.file_name, - row_position, score_getter(local_row_id)) + yield from _localized_candidates( + future.result(), context[0], context[1], deleted_positions) - # Match Java's separately bounded indexed and exact queues. The ANN - # queue must retain indexed_limit until exact reranking is complete. - indexed_candidates = _top_k(indexed_candidate_iter(), indexed_limit) - if self._configured_refine_factor(index_type) > 0 and indexed_candidates: - indexed_candidates = self._rerank_indexed( - plan, indexed_candidates, index_type) - else: - indexed_candidates = _top_k(indexed_candidates, self._limit) - exact_candidates = [] - if (self._table.options.vector_index_search_mode() - != GlobalIndexSearchMode.FAST): - exact_candidates = _top_k(self._raw_candidates(plan), self._limit) - candidates = _top_k( - (candidate - for group in (indexed_candidates, exact_candidates) - for candidate in group), - self._limit) - source_splits = [split.data_split for split in plan.splits()] - return PrimaryKeyScoredResult(plan.snapshot_id, source_splits, candidates) + return _top_k(indexed_candidate_iter(), indexed_limit) + + def _search_exact_candidates(self, plan): + return _top_k(self._raw_candidates(plan), self._limit) def _rerank_indexed(self, plan, candidates, index_type): source_splits = [split.data_split for split in plan.splits()] @@ -172,6 +175,15 @@ def _raw_candidates(self, plan): metric = self._search_metric( self._table.options.primary_key_vector_index_type( self._vector_column.name)) + for split, data_file, positions, vectors in self._raw_batches(plan): + for position, score in zip(positions, _iter_arrow_scores(vectors, self._query_vector, metric)): + if score is not None: + yield PrimaryKeySearchPosition( + _partition_bytes(split.data_split.partition), split.data_split.bucket, + data_file.file_name, position, score) + + def _raw_batches(self, plan): + """Read each selected file once, retaining its live physical row positions.""" read_table = self._table.copy_without_time_travel({ 'read.native.enabled': 'false', 'scan.native-plan.enabled': 'false'}) read_builder = read_table.new_read_builder().with_projection( @@ -199,18 +211,10 @@ def _raw_candidates(self, plan): if _allowed(split, data_file.file_name, position)) position_iter = iter(positions) for batch in reader.to_arrow([read_split]).to_batches(): - for score in _iter_arrow_scores(batch.column(0), self._query_vector, metric): - try: - row_position = next(position_iter) - except StopIteration: - raise ValueError( - "Raw vector read returned an unexpected row.") - if score is None: - continue - yield PrimaryKeySearchPosition( - _partition_bytes(split.data_split.partition), - split.data_split.bucket, data_file.file_name, - row_position, score) + batch_positions = list(islice(position_iter, batch.num_rows)) + if len(batch_positions) != batch.num_rows: + raise ValueError("Raw vector read returned an unexpected row.") + yield split, data_file, batch_positions, batch.column(0) try: next(position_iter) raise ValueError( @@ -219,6 +223,20 @@ def _raw_candidates(self, plan): pass +def _localized_candidates(result, split, source_meta, deleted_positions): + if result is None: + return + score_getter = result.score_getter() + for local_row_id in result.results(): + source, position = _localize(source_meta.source_files, local_row_id) + if position in deleted_positions.get((id(split), source.file_name), set()): + continue + if _allowed(split, source.file_name, position): + yield PrimaryKeySearchPosition( + _partition_bytes(split.data_split.partition), split.data_split.bucket, + source.file_name, position, score_getter(local_row_id)) + + def _rank_key(position): return ( -position.score, position.partition_bytes, position.bucket, diff --git a/paimon-python/pypaimon/execution/python/read/search/vector_search_read.py b/paimon-python/pypaimon/execution/python/read/search/vector_search_read.py index bd4fe301a6f1..738a3d98c86f 100644 --- a/paimon-python/pypaimon/execution/python/read/search/vector_search_read.py +++ b/paimon-python/pypaimon/execution/python/read/search/vector_search_read.py @@ -685,6 +685,16 @@ def __init__(self, table, limit, vector_column, query_vector, filter_=None, options=options) self._query_vector = query_vector + def _new_distributed_reader(self, executor): + if type(self) is not DataEvolutionVectorRead: + raise ValueError("Custom vector readers must define their own distributed execution.") + if not self._table.options.data_evolution_enabled(): + raise ValueError("Ray vector search supports only data-evolution tables for this reader.") + from pypaimon.execution.python.read.search.distributed_search import validate_query_vectors + from pypaimon.execution.python.read.search.distributed_vector_search import DistributedVectorSearchRead + validate_query_vectors([self._query_vector]) + return DistributedVectorSearchRead(self, executor) + def _read(self, splits, snapshot): self._index_metric = None index_splits, raw_splits = self._prepare_search_splits(splits, snapshot) @@ -732,6 +742,16 @@ def __init__(self, table, limit, vector_column, query_vectors, options=options) self._query_vectors = list(query_vectors) + def _new_distributed_reader(self, executor): + if type(self) is not BatchVectorSearchReadImpl: + raise ValueError("Custom vector readers must define their own distributed execution.") + if not self._table.options.data_evolution_enabled(): + raise ValueError("Ray vector search supports only data-evolution tables for this reader.") + from pypaimon.execution.python.read.search.distributed_search import validate_query_vectors + from pypaimon.execution.python.read.search.distributed_vector_search import DistributedBatchVectorSearchRead + validate_query_vectors(self._query_vectors) + return DistributedBatchVectorSearchRead(self, executor) + def _read_batch(self, splits, snapshot): self._index_metric = None n = len(self._query_vectors) diff --git a/paimon-python/pypaimon/multimodal/query.py b/paimon-python/pypaimon/multimodal/query.py index e194628b33c8..20840d98fecf 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -675,7 +675,8 @@ def __init__( def to_arrow(self, *, execution="local", concurrency=None, ray_remote_args=None): """Execute a vector query locally or on Ray workers, returning an Arrow table. - Ray execution supports data-evolution tables. ``concurrency`` bounds + Ray execution supports data-evolution tables and vector search on + primary-key tables using the Python read backend. ``concurrency`` bounds in-flight tasks (defaults to 4); ``ray_remote_args`` configures their resources and retries. Result lookup runs on the driver. """ @@ -856,9 +857,10 @@ def __init__( def to_arrow(self, *, execution="local", concurrency=None, ray_remote_args=None): """Return one Arrow table per query, optionally searching splits on Ray. - Ray execution supports data-evolution tables. ``concurrency`` bounds + Ray execution supports data-evolution tables and vector search on + primary-key tables using the Python read backend. ``concurrency`` bounds in-flight tasks (defaults to 4); ``ray_remote_args`` configures their - resources and retries. Shared result lookup runs on the driver, + resources and retries. Result lookup runs on the driver, using the same snapshot as all workers, including refinement tasks. """ if execution == "local": @@ -899,10 +901,12 @@ def to_ray(self, *, execution="ray", concurrency=None, ray_remote_args=None, def _read_batch_results(self, results): from pypaimon.globalindex.global_index_result import GlobalIndexResult + from pypaimon.table.source.global_index_split_result import GlobalIndexSplitResult from pypaimon.utils.roaring_bitmap import RoaringBitmap64 projection = self._effective_projection() - if (len(results) <= 1 or self._metadata_only_result() + if (len(results) <= 1 or any(isinstance(result, GlobalIndexSplitResult) for result in results) + or self._metadata_only_result() or (projection and not self._configured_read_builder()._output_fields())): return [self._read_global_index_result(result) for result in results] diff --git a/paimon-python/pypaimon/tests/distributed_reader_selection_test.py b/paimon-python/pypaimon/tests/distributed_reader_selection_test.py new file mode 100644 index 000000000000..cbcd73989226 --- /dev/null +++ b/paimon-python/pypaimon/tests/distributed_reader_selection_test.py @@ -0,0 +1,178 @@ +# 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. + +"""Distributed vector reader selection without a Ray runtime or index files.""" + +from types import SimpleNamespace + +import pytest + +from pypaimon.common.options.core_options import CoreOptions +from pypaimon.common.options.options import Options +from pypaimon.execution.python.read.search.distributed_primary_key_vector_search import ( + DistributedBatchPrimaryKeyVectorRead, DistributedPrimaryKeyVectorRead, +) +from pypaimon.execution.python.read.search.distributed_search import create_distributed_reader +from pypaimon.execution.python.read.search.distributed_vector_search import ( + DistributedBatchVectorSearchRead, DistributedVectorSearchRead, +) +from pypaimon.execution.python.read.search.primary_key_batch_vector_read import BatchPrimaryKeyVectorRead +from pypaimon.execution.python.read.search.primary_key_vector_read import PrimaryKeyVectorRead +from pypaimon.execution.python.read.search.vector_search_read import ( + BatchVectorSearchReadImpl, DataEvolutionVectorRead, +) +from pypaimon.table.source.search_types import PrimaryKeyVectorScanPlan, VectorSearchScanPlan + + +VECTOR_READERS = [ + DataEvolutionVectorRead, BatchVectorSearchReadImpl, + PrimaryKeyVectorRead, BatchPrimaryKeyVectorRead, +] +SUPPORTED_READERS = [ + (DataEvolutionVectorRead, DistributedVectorSearchRead), + (BatchVectorSearchReadImpl, DistributedBatchVectorSearchRead), + (PrimaryKeyVectorRead, DistributedPrimaryKeyVectorRead), + (BatchPrimaryKeyVectorRead, DistributedBatchPrimaryKeyVectorRead), +] + + +class RecordingExecutor: + """Expose query views and fail if selection tries to schedule work.""" + + def __init__(self, concurrency=3, parent=None): + self.concurrency = concurrency + self.views = [] if parent is None else parent.views + self.calls = {"map": 0} if parent is None else parent.calls + + def with_concurrency(self, concurrency): + view = RecordingExecutor(concurrency, self) + self.views.append(view) + return view + + def map(self, tasks, ordered=False): + self.calls["map"] += 1 + raise AssertionError("Reader selection must not schedule tasks.") + + +def _reader(reader_type, evolution=None): + primary_key = issubclass(reader_type, PrimaryKeyVectorRead) + evolution = not primary_key if evolution is None else evolution + options = {"data-evolution.enabled": str(evolution).lower()} + table = SimpleNamespace( + is_primary_key_table=primary_key, + options=CoreOptions(Options(options)), + table_schema=SimpleNamespace(options=options), + ) + column = SimpleNamespace(name="embedding") + queries = ([[1., 0.], [0., 1.]] if issubclass( + reader_type, (BatchVectorSearchReadImpl, BatchPrimaryKeyVectorRead)) else [1., 0.]) + return reader_type(table, 2, column, queries, options={"ivf.refine_factor": "2"}) + + +@pytest.mark.parametrize("reader_type,distributed_type", SUPPORTED_READERS) +def test_supported_reader_selects_its_distributed_execution(reader_type, distributed_type): + reader = _reader(reader_type) + executor = RecordingExecutor() + + distributed = create_distributed_reader(reader, executor) + + assert type(distributed) is distributed_type + assert distributed._table is reader._table + assert distributed._limit == reader._limit + assert distributed._executor is executor.views[0] + assert distributed._executor.concurrency == executor.concurrency + assert executor.calls["map"] == 0 + + +@pytest.mark.parametrize("reader_type", [ + DataEvolutionVectorRead, BatchVectorSearchReadImpl, +]) +def test_data_evolution_reader_rejects_plain_append_table(reader_type): + executor = RecordingExecutor() + + with pytest.raises((ValueError, NotImplementedError)): + create_distributed_reader(_reader(reader_type, evolution=False), executor) + + assert executor.calls["map"] == 0 + + +@pytest.mark.parametrize("reader_type", VECTOR_READERS) +@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")]) +def test_nonfinite_queries_fail_before_task_submission(reader_type, value): + reader = _reader(reader_type) + executor = RecordingExecutor() + if hasattr(reader, "_query_vectors"): + # Include a valid first query so every batch entry must be validated. + reader._query_vectors[1][0] = value + read = reader.read_batch_plan + else: + reader._query_vector[0] = value + read = reader.read_plan + plan = (PrimaryKeyVectorScanPlan(0, []) if reader._table.is_primary_key_table + else VectorSearchScanPlan([])) + + with pytest.raises(ValueError, match="finite"): + read(plan, executor=executor) + + assert executor.calls["map"] == 0 + + +@pytest.mark.parametrize("reader_type,_", SUPPORTED_READERS) +def test_custom_subclass_must_explicitly_opt_into_distributed_execution(reader_type, _): + class CustomReader(reader_type): + pass + + executor = RecordingExecutor() + + with pytest.raises((ValueError, NotImplementedError)): + create_distributed_reader(_reader(CustomReader), executor) + + assert executor.calls["map"] == 0 + + +def test_custom_reader_hook_preserves_its_execution_semantics(): + result = object() + + class CustomReader(DataEvolutionVectorRead): + def _new_distributed_reader(self, executor): + self.selected_executor = executor + return self + + def _read(self, splits, snapshot): + return result + + reader = _reader(CustomReader) + executor = RecordingExecutor() + + assert reader.read_plan(VectorSearchScanPlan([]), executor=executor) is result + assert reader.selected_executor is executor.views[0] + assert executor.calls["map"] == 0 + + +def test_reusing_executor_creates_independent_query_views(): + reader = _reader(PrimaryKeyVectorRead) + executor = RecordingExecutor() + + first = create_distributed_reader(reader, executor) + second = create_distributed_reader(reader, executor) + + assert first._executor is executor.views[0] + assert second._executor is executor.views[1] + assert first._executor is not second._executor + assert first._executor is not executor + assert second._executor is not executor + assert executor.calls["map"] == 0 diff --git a/paimon-python/pypaimon/tests/pk_search_capabilities_test.py b/paimon-python/pypaimon/tests/pk_search_capabilities_test.py new file mode 100644 index 000000000000..a7cd619b478f --- /dev/null +++ b/paimon-python/pypaimon/tests/pk_search_capabilities_test.py @@ -0,0 +1,80 @@ +# 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 tarfile +from types import SimpleNamespace + +import pytest + +from pypaimon.catalog.catalog_factory import CatalogFactory +from pypaimon.multimodal.query import BatchVectorQuery, VectorQuery +from pypaimon.table.source.primary_key_scored_result import PrimaryKeyScoredResult +from pypaimon.tests.primary_key_global_index_golden_test import GOLDEN_ARCHIVE + + +def _query(query_type): + table = SimpleNamespace( + fields=[], options=SimpleNamespace(data_evolution_enabled=lambda: False)) + vectors = [[1., 0.]] if query_type is BatchVectorQuery else [1., 0.] + return query_type(table, vectors, "embedding") + + +@pytest.mark.parametrize("query_type", [VectorQuery, BatchVectorQuery]) +def test_pk_ray_dataset_conversion_requires_row_id_results(query_type): + pytest.importorskip("ray") + from pypaimon.ray.search_result import read_search_result + + query = _query(query_type) + with pytest.raises(NotImplementedError, match="require a data-evolution table"): + read_search_result(query, PrimaryKeyScoredResult(0, [], []), 1, {}, None) + + +@pytest.mark.parametrize("query_type", [VectorQuery, BatchVectorQuery]) +@pytest.mark.parametrize("method", ["with_score", "order_by_score"]) +def test_pk_score_outputs_rejected_when_configuring_query(query_type, method): + query = _query(query_type) + with pytest.raises(NotImplementedError, match="Search score output and ordering require a data-evolution"): + getattr(query, method)() + + +@pytest.fixture +def pk_table(tmp_path): + pytest.importorskip("paimon_vindex") + with tarfile.open(GOLDEN_ARCHIVE, "r:gz") as archive: + archive.extractall(tmp_path) + return CatalogFactory.create({"warehouse": str(tmp_path)}).get_table( + "default.test_pk_vector_golden").copy({ + "read.native.enabled": "false", "scan.native-plan.enabled": "false", + "vector-index.search-mode": "full"}) + + +@pytest.mark.parametrize("projection", [None, [], ["id"], ["id", "id"], ["missing"]]) +@pytest.mark.parametrize("post_filter", [None, "id >= 3"]) +def test_pk_batch_physical_lookup_matches_single_queries(pk_table, projection, post_filter): + vectors = [[1., 0., 0., 0.], [0., 1., 0., 0.], [1., 0., 0., 0.]] + + def configure(query): + if projection is not None: + query.select(projection) + if post_filter is not None: + query.where(post_filter) + return query.limit(3) + + expected = [configure(VectorQuery(pk_table, vector, "embedding")).to_arrow() for vector in vectors] + actual = configure(BatchVectorQuery(pk_table, vectors, "embedding")).to_arrow() + assert len(actual) == len(expected) + assert all(batch.equals(single) for batch, single in zip(actual, expected)) diff --git a/paimon-python/pypaimon/tests/ray_primary_key_vector_search_test.py b/paimon-python/pypaimon/tests/ray_primary_key_vector_search_test.py new file mode 100644 index 000000000000..893035f95e15 --- /dev/null +++ b/paimon-python/pypaimon/tests/ray_primary_key_vector_search_test.py @@ -0,0 +1,645 @@ +# 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 os +import tarfile +from dataclasses import replace +from pathlib import Path +from unittest.mock import Mock, patch + +import pytest + +ray = pytest.importorskip("ray") + +from pypaimon.catalog.catalog_factory import CatalogFactory +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 +from pypaimon.multimodal.query import VectorQuery +from pypaimon.ray.search_execution import RaySearchTaskExecutor +from pypaimon.table.source.search_types import PrimaryKeyVectorScanPlan +from pypaimon.tests.primary_key_global_index_golden_test import GOLDEN_ARCHIVE, _read_search_result +from pypaimon.utils.range import Range + + +QUERIES = [[1., 0., 0., 0.], [0., 1., 0., 0.], [1., 0., 0., 0.]] + + +@pytest.fixture(scope="module") +def ray_cluster(): + started = not ray.is_initialized() + if started: + ray.init(address="local", num_cpus=2, include_dashboard=False, + object_store_memory=100 * 1024 * 1024) + yield + if started: + ray.shutdown() + + +@pytest.fixture +def table(tmp_path): + with tarfile.open(GOLDEN_ARCHIVE, "r:gz") as archive: + archive.extractall(tmp_path) + return CatalogFactory.create({"warehouse": str(tmp_path)}).get_table( + "default.test_pk_vector_golden").copy({ + "read.native.enabled": "false", "scan.native-plan.enabled": "false", + "vector-index.search-mode": "full"}) + + +def builder(table, limit=2, options=None): + return (table.new_vector_search_builder().with_vector_column("embedding") + .with_query_vector([1., 0., 0., 0.]).with_limit(limit) + .with_options(options or {})) + + +@pytest.mark.parametrize("refine", [False, True]) +@pytest.mark.parametrize("pre_filter", [None, "id >= 3", "id = 999"]) +def test_indexed_query_matches_local(table, ray_cluster, refine, pre_filter): + pytest.importorskip("paimon_vindex") + query = VectorQuery(table, [1., 0., 0., 0.], "embedding", + vector_options={"ivf.refine_factor": "2"} if refine else {}, + pre_filter=pre_filter).limit(3) + expected = query.to_arrow().to_pylist() + assert query.to_arrow(execution="ray", concurrency=2).to_pylist() == expected + + +@pytest.mark.parametrize("mode", ["full", "fast"]) +@pytest.mark.parametrize("batch", [False, True]) +def test_raw_plan_preserves_filters_and_search_mode(table, ray_cluster, mode, batch): + read_table = table.copy({"vector-index.search-mode": mode}) + request = (batch_builder if batch else builder)(read_table, limit=10) + original = request.new_vector_search_scan().scan() + splits = [replace(split, payloads=(), uncovered_data_files=tuple( + file.file_name for file in split.data_split.files), + row_ranges_by_file={file.file_name: (Range(1, 2),) + for file in split.data_split.files}) for split in original.splits()] + plan = PrimaryKeyVectorScanPlan(original.snapshot_id, splits, original.snapshot()) + reader = request.new_batch_vector_search_read() if batch else request.new_vector_search_read() + read = reader.read_batch_plan if batch else reader.read_plan + expected, actual = read(plan), read(plan, executor=RaySearchTaskExecutor(concurrency=2)) + for local, distributed in zip(expected if batch else [expected], actual if batch else [actual]): + assert distributed.positions == local.positions + assert len(distributed.positions) == (0 if mode == "fast" else 2) + assert _read_search_result(read_table, distributed).to_pylist() == _read_search_result( + read_table, local).to_pylist() + + +def test_empty_plan_never_dispatches(table): + request = builder(table) + executor = RaySearchTaskExecutor(concurrency=2) + with patch.object(executor, "map", side_effect=AssertionError("Unexpected tasks")): + # with_concurrency creates a query view, so exercise the reader directly. + from pypaimon.execution.python.read.search.distributed_primary_key_vector_search import ( + DistributedPrimaryKeyVectorRead, + ) + reader = DistributedPrimaryKeyVectorRead(request.new_vector_search_read(), executor) + result = reader.read_plan(PrimaryKeyVectorScanPlan(0, [])) + assert result.snapshot_id == 0 + assert not result.positions + + +@pytest.mark.parametrize("value", [float("nan"), float("inf")]) +def test_rejects_nonfinite_query_before_dispatch(table, value): + request = builder(table).with_query_vector([value, 0., 0., 0.]) + with pytest.raises(ValueError, match="finite query"): + request.new_vector_search_read().read_plan( + request.new_vector_search_scan().scan(), executor=RaySearchTaskExecutor()) + + +class InlineExecutor: + concurrency = 2 + + def with_concurrency(self, concurrency): + return self + + def map(self, tasks, ordered=False): + for ordinal, task in reversed(list(enumerate(tasks))): + yield ordinal, task.execute() + + +def two_bucket_plan(table): + """Exercise supplied bucket plans with real Java-generated index/data files.""" + import copy + import shutil + + plan = builder(table).new_vector_search_scan().scan() + first = plan.splits()[0] + second_data = copy.deepcopy(first.data_split) + second_data._bucket = 1 + destination = Path(table.table_path) / "bucket-1" + destination.mkdir() + for data_file in second_data.files: + target = destination / data_file.file_name + shutil.copyfile(data_file.file_path, target) + data_file.file_path = str(target) + second = replace(first, data_split=second_data) + return PrimaryKeyVectorScanPlan(plan.snapshot_id, [first, second], plan.snapshot()) + + +def _report_task(context, item): + import time + + operation, original_item, directory = item + split = original_item[1] + Path(directory, str(split.data_split.bucket)).write_text(str(os.getpid())) + deadline = time.monotonic() + 15 + # Both tasks must be alive together, rather than one process reusing the worker. + while len(list(Path(directory).iterdir())) < 2: + if time.monotonic() > deadline: + raise RuntimeError("Second bucket task did not start") + time.sleep(0.01) + return os.getpid(), operation(context, original_item) + + +class ReportingExecutor: + concurrency = 2 + + def __init__(self, directory): + self.directory = str(directory) + self.pids = set() + self.delegate = RaySearchTaskExecutor(concurrency=2, ray_remote_args={"num_cpus": 1}) + + def with_concurrency(self, concurrency): + return self + + def map(self, tasks, ordered=False): + from contextlib import closing + work = [replace(task, operation=_report_task, + item=(task.operation, task.item, self.directory)) for task in tasks] + with closing(self.delegate.map(work, ordered=ordered)) as results: + for ordinal, (pid, result) in results: + self.pids.add(pid) + yield ordinal, result + + +@pytest.mark.parametrize("raw", [False, True]) +def test_two_buckets_execute_on_distinct_workers(table, ray_cluster, tmp_path, raw): + pytest.importorskip("paimon_vindex") + plan = two_bucket_plan(table) + if raw: + plan = PrimaryKeyVectorScanPlan(plan.snapshot_id, [ + replace(split, payloads=(), uncovered_data_files=tuple( + file.file_name for file in split.data_split.files)) for split in plan.splits() + ], plan.snapshot()) + request = builder(table, limit=5) + reader = request.new_vector_search_read() + expected = reader.read_plan(plan) + directory = tmp_path / "workers" + directory.mkdir() + executor = ReportingExecutor(directory) + actual = reader.read_plan(plan, executor=executor) + assert actual.positions == expected.positions + assert {position.bucket for position in actual.positions} == {0, 1} + assert len(executor.pids) == 2 + assert os.getpid() not in executor.pids + assert _read_search_result(table, actual).to_pylist() == _read_search_result( + table, expected).to_pylist() + + +def test_global_ann_budget_is_preserved_before_refinement(table): + from concurrent.futures import Future + from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult + + plan = two_bucket_plan(table) + second = plan.splits()[1] + plan = PrimaryKeyVectorScanPlan(plan.snapshot_id, [plan.splits()[0], replace( + second, payloads=(replace(second.payloads[0], file_name="bucket-1-index"),))], plan.snapshot()) + request = builder(table, limit=1, options={"ivf.refine_factor": "2"}) + seen = [] + + def evaluate(reader, start, end, files, query, limit, include): + seen.append(limit) + reader._set_index_metric("l2") + future = Future() + future.set_result(DictBasedScoredIndexResult( + {0: 98., 1: 97.} if files[0].file_name == "bucket-1-index" else {0: 100., 1: 99.})) + return future + + def refine(reader, full_plan, candidates, index_type): + # The global ANN budget can contain multiple candidates from one bucket. + assert [(p.bucket, p.row_position) for p in candidates] == [(0, 0), (0, 1)] + return [candidates[1].with_score(200.)] + + with patch.object(PrimaryKeyVectorRead, "_eval", evaluate), patch.object( + PrimaryKeyVectorRead, "_rerank_indexed", refine): + result = request.new_vector_search_read().read_plan(plan, executor=InlineExecutor()) + assert seen == [2, 2] + assert result.positions[0].row_position == 1 + assert result.positions[0].score == 200. + + +@pytest.mark.parametrize("batch", [False, True]) +def test_metric_mismatch_on_empty_bucket_closes_pending_results(table, batch): + from pypaimon.execution.python.read.search.distributed_primary_key_vector_search import ( + DistributedPrimaryKeyVectorRead, DistributedBatchPrimaryKeyVectorRead, + ) + closed = [] + + class ConflictingExecutor: + def map(self, tasks): + assert len(tasks) == 2 + try: + yield 0, ("l2", [[] for _ in QUERIES] if batch else []) + yield 1, ("cosine", [[] for _ in QUERIES] if batch else []) + finally: + closed.append(True) + + if batch: + reader = DistributedBatchPrimaryKeyVectorRead( + batch_builder(table).new_batch_vector_search_read(), ConflictingExecutor()) + else: + reader = DistributedPrimaryKeyVectorRead(builder(table).new_vector_search_read(), ConflictingExecutor()) + with pytest.raises(ValueError, match="different metrics"): + (reader.read_batch_plan if batch else reader.read_plan)(two_bucket_plan(table)) + assert closed == [True] + + +@pytest.mark.parametrize("indexed", [False, True]) +def test_deletion_vectors_filter_worker_candidates(table, ray_cluster, tmp_path, indexed): + from pypaimon.execution.python.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector + from pypaimon.table.source.deletion_file import DeletionFile + + original = builder(table).new_vector_search_scan().scan() + vector = BitmapDeletionVector() + vector.delete(0) + data = vector.serialize() + path = tmp_path / "deleted.dv" + path.write_bytes(data) + deletion = DeletionFile(str(path), 0, len(data) - 8, 1) + split = original.splits()[0] + import copy + source = copy.deepcopy(split.data_split) + source.data_deletion_files = [deletion] + selected = replace(split, data_split=source) + if not indexed: + selected = replace(selected, payloads=(), uncovered_data_files=tuple( + file.file_name for file in source.files)) + plan = PrimaryKeyVectorScanPlan(original.snapshot_id, [selected], original.snapshot()) + reader = builder(table, limit=10).new_vector_search_read() + expected = reader.read_plan(plan) + actual = reader.read_plan(plan, executor=RaySearchTaskExecutor(concurrency=2)) + assert actual.positions == expected.positions + assert len(actual.positions) == 3 + assert all(position.row_position != 0 for position in actual.positions) + + +@pytest.mark.parametrize("batch", [False, True]) +def test_worker_failure_propagates_and_reader_can_be_reused(table, ray_cluster, batch): + original = builder(table).new_vector_search_scan().scan() + split = original.splits()[0] + missing = replace(split.payloads[0], file_name="missing-vector-index") + plan = PrimaryKeyVectorScanPlan( + original.snapshot_id, [replace(split, payloads=(missing,))], original.snapshot()) + reader = (batch_builder(table).new_batch_vector_search_read() if batch + else builder(table).new_vector_search_read()) + read = reader.read_batch_plan if batch else reader.read_plan + with pytest.raises(Exception, match="missing-vector-index"): + read(plan, executor=RaySearchTaskExecutor(concurrency=2)) + # No failed execution state is retained on the original reader. + expected = read(original) + actual = read(original, executor=RaySearchTaskExecutor()) + assert ([result.positions for result in actual] if batch else actual.positions) == ( + [result.positions for result in expected] if batch else expected.positions) + + +@pytest.mark.parametrize("refine", [False, True]) +@pytest.mark.parametrize("batch", [False, True]) +def test_mixed_index_and_raw_buckets_share_persisted_metric(table, ray_cluster, refine, batch): + table = table.copy({"fields.embedding.distance.metric": "cosine"}) + original = two_bucket_plan(table) + first, second = original.splits() + second = replace(second, payloads=(), uncovered_data_files=tuple( + file.file_name for file in second.data_split.files)) + plan = PrimaryKeyVectorScanPlan(original.snapshot_id, [first, second], original.snapshot()) + request = (batch_builder if batch else builder)( + table, limit=6, options={"ivf.refine_factor": "2"} if refine else {}) + reader = request.new_batch_vector_search_read() if batch else request.new_vector_search_read() + read = reader.read_batch_plan if batch else reader.read_plan + expected, actual = read(plan), read(plan, executor=RaySearchTaskExecutor(concurrency=2)) + for local, distributed in zip(expected if batch else [expected], actual if batch else [actual]): + assert distributed.positions == local.positions + assert {position.bucket for position in distributed.positions} == {0, 1} + + +def _retry_at_snapshot(context, item): + from pypaimon.snapshot.time_travel_util import TimeTravelUtil + + operation, original_item, marker, snapshot_id = item + assert TimeTravelUtil.resolve_snapshot(context[0]).id == snapshot_id + try: + fd = os.open(marker, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + except FileExistsError: + return operation(context, original_item) + os.close(fd) + raise RuntimeError("Retry the pinned bucket search") + + +@pytest.mark.parametrize("batch", [False, True]) +def test_task_retry_does_not_resolve_latest_snapshot(table, ray_cluster, tmp_path, batch): + import pyarrow as pa + from contextlib import closing + from pypaimon.snapshot.time_travel_util import TimeTravelUtil + + request = (batch_builder if batch else builder)(table, options={"ivf.refine_factor": "2"}) + plan = request.new_vector_search_scan().scan() + reader = request.new_batch_vector_search_read() if batch else request.new_vector_search_read() + read = reader.read_batch_plan if batch else reader.read_plan + expected = read(plan) + write_builder = table.new_batch_write_builder() + writer, commit = write_builder.new_write(), write_builder.new_commit() + try: + writer.write_arrow(pa.table({"id": pa.array([99], type=pa.int32()), + "embedding": pa.array([[1., 0., 0., 0.]], + type=pa.list_(pa.float32(), 4))})) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + assert TimeTravelUtil.resolve_snapshot(table).id > plan.snapshot_id + + class RetryingExecutor: + concurrency = 2 + + def with_concurrency(self, concurrency): + return self + + def map(self, tasks, ordered=False): + delegate = RaySearchTaskExecutor(concurrency=2, ray_remote_args={ + "max_retries": 1, "retry_exceptions": True}) + work = [replace(task, operation=_retry_at_snapshot, + item=(task.operation, task.item, str(tmp_path / "retry"), plan.snapshot_id)) + for task in tasks] + with closing(delegate.map(work, ordered=ordered)) as results: + yield from results + + actual = read(plan, executor=RetryingExecutor()) + for local, distributed in zip(expected if batch else [expected], actual if batch else [actual]): + assert distributed.snapshot_id == plan.snapshot_id + assert distributed.positions == local.positions + + +def batch_builder(table, queries=None, limit=2, options=None): + return (table.new_batch_vector_search_builder().with_vector_column("embedding") + .with_query_vectors(QUERIES if queries is None else queries).with_limit(limit) + .with_options(options or {})) + + +@pytest.mark.parametrize("refine", [False, True]) +@pytest.mark.parametrize("pre_filter", [None, "id >= 3", "id = 999"]) +def test_batch_public_api_matches_independent_queries(table, ray_cluster, refine, pre_filter): + from pypaimon.multimodal.table import MultimodalTable + + options = {"ivf.refine_factor": "2"} if refine else {} + expected = [VectorQuery(table, query, "embedding", vector_options=options, + pre_filter=pre_filter).select(["id"]).limit(3).to_arrow().to_pylist() + for query in QUERIES] + view = MultimodalTable(None, table.identifier.get_full_name(), table) + batch = view.search_vectors(QUERIES, column="embedding", options=options, + pre_filter=pre_filter).select(["id"]).limit(3) + assert batch.to_list() == expected + results = batch.to_arrow(execution="ray", concurrency=2) + assert [result.to_pylist() for result in results] == expected + assert all(result.column_names == ["id"] for result in results) + assert expected[0] == expected[2] + + +@pytest.mark.parametrize("raw", [False, True]) +def test_batch_two_buckets_share_tasks_across_queries(table, ray_cluster, tmp_path, raw): + plan = two_bucket_plan(table) + if raw: + plan = PrimaryKeyVectorScanPlan(plan.snapshot_id, [replace( + split, payloads=(), uncovered_data_files=tuple(file.file_name for file in split.data_split.files), + row_ranges_by_file={file.file_name: (Range(1, 2),) for file in split.data_split.files}) + for split in plan.splits()], plan.snapshot()) + request = batch_builder(table, limit=5) + reader = request.new_batch_vector_search_read() + expected = [] + for query in QUERIES: + single = builder(table, limit=5).with_query_vector(query).new_vector_search_read() + expected.append(single.read_plan(plan).positions) + directory = tmp_path / "batch-workers" + directory.mkdir() + executor = ReportingExecutor(directory) + results = reader.read_batch_plan(plan, executor=executor) + assert [result.positions for result in results] == expected + assert len(executor.pids) == 2 + assert os.getpid() not in executor.pids + + +def test_batch_reuses_index_reader_across_query_blocks_and_closes_on_failure(table): + request = batch_builder(table, queries=QUERIES * 44) + plan = request.new_vector_search_scan().scan() + reader = request.new_batch_vector_search_read() + opened = [] + original = reader._open_offset_reader + + def open_reader(*args, **kwargs): + index, offset = original(*args, **kwargs) + index.close = Mock(wraps=index.close) + offset.visit_batch_vector_search = Mock(wraps=offset.visit_batch_vector_search) + opened.append((index, offset)) + return index, offset + + with patch.object(reader, "_open_offset_reader", side_effect=open_reader): + results = reader.read_batch_plan(plan) + assert len(results) == 132 + assert len(opened) == sum(len(split.payloads) for split in plan.splits()) + for index, offset in opened: + assert offset.visit_batch_vector_search.call_count == 3 + index.close.assert_called_once() + + index, offset = original([plan.splits()[0].payloads[0]], 0, 3) + with patch.object(index, "close", wraps=index.close) as close, \ + patch.object(offset, "visit_batch_vector_search", side_effect=RuntimeError("batch search failed")), \ + patch.object(reader, "_open_offset_reader", return_value=(index, offset)): + with pytest.raises(RuntimeError, match="batch search failed"): + reader.read_batch_plan(plan) + close.assert_called_once() + + +def test_batch_refinement_reads_candidate_union_once(table): + reader = batch_builder(table, limit=3, options={"ivf.refine_factor": "2"}).new_batch_vector_search_read() + plan = builder(table).new_vector_search_scan().scan() + original = reader._raw_batches + read_files = [] + + def raw_batches(selected): + for split, file, positions, vectors in original(selected): + read_files.append((file.file_name, tuple(positions))) + yield split, file, positions, vectors + + with patch.object(reader, "_raw_batches", side_effect=raw_batches): + actual = reader.read_batch_plan(plan) + expected = [builder(table, limit=3, options={"ivf.refine_factor": "2"}) + .with_query_vector(query).new_vector_search_read().read_plan(plan).positions for query in QUERIES] + assert [result.positions for result in actual] == expected + assert len(read_files) == 1 + + +def test_batch_keeps_each_query_global_refinement_budget(table): + from pypaimon.table.source.primary_key_scored_result import PrimaryKeySearchPosition, _partition_bytes + + plan = two_bucket_plan(table) + reader = batch_builder(table, queries=QUERIES[:2], limit=1, + options={"ivf.refine_factor": "2"}).new_batch_vector_search_read() + + def search(worker, local_plan, limit): + assert limit == 2 + worker._set_index_metric("l2") + split = local_plan.splits()[0].data_split + return [[PrimaryKeySearchPosition(_partition_bytes(split.partition), split.bucket, + split.files[0].file_name, position, + 100. - position - (10. if split.bucket != i else 0.)) + for position in (0, 1)] for i in range(2)] + + def refine(worker, full_plan, candidates, index_type): + assert [[(p.bucket, p.row_position) for p in group] for group in candidates] == [ + [(0, 0), (0, 1)], [(1, 0), (1, 1)]] + return [[group[1].with_score(200.)] for group in candidates] + + with patch.object(BatchPrimaryKeyVectorRead, "_search_indexed_candidates", search), \ + patch.object(BatchPrimaryKeyVectorRead, "_rerank_indexed_batch", refine): + actual = reader.read_batch_plan(plan, executor=InlineExecutor()) + assert [(result.positions[0].bucket, result.positions[0].row_position) for result in actual] == [(0, 1), (1, 1)] + + +@pytest.mark.parametrize("indexed", [False, True]) +def test_batch_deletion_vectors_and_physical_ranges(table, ray_cluster, tmp_path, indexed): + import copy + from pypaimon.execution.python.deletionvectors.bitmap_deletion_vector import BitmapDeletionVector + from pypaimon.table.source.deletion_file import DeletionFile + + plan = two_bucket_plan(table) + vector = BitmapDeletionVector() + vector.delete(1) + data = vector.serialize() + path = tmp_path / "batch-deleted.dv" + path.write_bytes(data) + splits = [] + for split in plan.splits(): + source = copy.deepcopy(split.data_split) + source.data_deletion_files = [DeletionFile(str(path), 0, len(data) - 8, 1)] + splits.append(replace(split, data_split=source, + payloads=split.payloads if indexed else (), + uncovered_data_files=() if indexed else tuple(f.file_name for f in source.files), + row_ranges_by_file={f.file_name: (Range(1, 2),) for f in source.files})) + plan = PrimaryKeyVectorScanPlan(plan.snapshot_id, splits, plan.snapshot()) + reader = batch_builder(table, limit=10).new_batch_vector_search_read() + actual = reader.read_batch_plan(plan, executor=RaySearchTaskExecutor(concurrency=2)) + expected = [builder(table, limit=10).with_query_vector(query).new_vector_search_read().read_plan(plan).positions + for query in QUERIES] + assert [result.positions for result in actual] == expected + assert all(len(result.positions) == 2 for result in actual) + assert all(position.row_position == 2 for result in actual for position in result.positions) + + +@pytest.mark.parametrize("queries", [[], [[1.], [1., 2.]], [[float("nan")]], [[float("inf")]]]) +def test_batch_rejects_invalid_queries_before_dispatch(table, queries): + request = batch_builder(table, queries=queries) + with pytest.raises(ValueError): + request.new_batch_vector_search_read() + + +def test_batch_empty_plan_preserves_result_count_without_tasks(table): + reader = batch_builder(table).new_batch_vector_search_read() + executor = InlineExecutor() + with patch.object(executor, "map", side_effect=AssertionError("Unexpected tasks")): + results = reader.read_batch_plan(PrimaryKeyVectorScanPlan(0, []), executor=executor) + assert len(results) == len(QUERIES) + assert all(not result.positions and result.snapshot_id == 0 for result in results) + + +def test_batch_where_and_pandas_preserve_per_query_limits(table, ray_cluster): + from pypaimon.multimodal.query import BatchVectorQuery + + query = BatchVectorQuery(table, QUERIES, "embedding").select(["id"]).where("id >= 3").limit(2) + expected = [VectorQuery(table, vector, "embedding").select(["id"]).where("id >= 3") + .limit(2).to_arrow().to_pylist() for vector in QUERIES] + assert query.to_list(execution="ray", concurrency=2) == expected + assert [frame.to_dict("records") for frame in query.to_pandas()] == expected + + +def test_batch_raw_scan_reads_each_file_once(table): + plan = builder(table).new_vector_search_scan().scan() + plan = PrimaryKeyVectorScanPlan(plan.snapshot_id, [replace( + split, payloads=(), uncovered_data_files=tuple(f.file_name for f in split.data_split.files)) + for split in plan.splits()], plan.snapshot()) + reader = batch_builder(table).new_batch_vector_search_read() + read_type = type(table.new_read_builder().new_read()) + original = read_type.to_arrow + reads = [] + + def read(execution, splits, *args, **kwargs): + reads.extend(file.file_name for split in splits for file in split.files) + return original(execution, splits, *args, **kwargs) + + with patch.object(read_type, "to_arrow", read): + actual = reader.read_batch_plan(plan) + assert len(actual) == len(QUERIES) + assert reads == [file.file_name for split in plan.splits() for file in split.data_split.files] + + +def test_batch_wrong_result_arity_closes_scheduled_work(table): + from pypaimon.execution.python.read.search.distributed_primary_key_vector_search import ( + DistributedBatchPrimaryKeyVectorRead, + ) + closed = [] + + class InvalidExecutor: + def map(self, tasks): + try: + yield 0, ("l2", [[]]) + finally: + closed.append(True) + + reader = DistributedBatchPrimaryKeyVectorRead( + batch_builder(table).new_batch_vector_search_read(), InvalidExecutor()) + with pytest.raises(ValueError, match="wrong query count"): + reader.read_batch_plan(two_bucket_plan(table)) + assert closed == [True] + + +@pytest.mark.parametrize("projection", [None, [], ["id"], ["embedding"], ["id", "id"], ["missing"]]) +@pytest.mark.parametrize("queries", [QUERIES[:1], QUERIES]) +def test_batch_physical_lookup_preserves_projections(table, ray_cluster, projection, queries): + from pypaimon.multimodal.query import BatchVectorQuery + + batch = BatchVectorQuery(table, queries, "embedding").limit(3) + singles = [VectorQuery(table, query, "embedding").limit(3) for query in queries] + if projection is not None: + batch.select(projection) + for single in singles: + single.select(projection) + expected = [single.to_arrow() for single in singles] + actual = batch.to_arrow(execution="ray", concurrency=2) + assert len(actual) == len(queries) + assert all(result.equals(local) for result, local in zip(actual, expected)) + + +def test_batch_ray_dataset_output_remains_unsupported(table, ray_cluster): + from pypaimon.multimodal.query import BatchVectorQuery + + with pytest.raises(NotImplementedError, match="data-evolution"): + BatchVectorQuery(table, QUERIES, "embedding").to_ray(execution="local") + + +def test_batch_split_only_read_requires_captured_snapshot(table): + request = batch_builder(table) + reader = request.new_batch_vector_search_read() + plan = request.new_vector_search_scan().scan() + with pytest.raises(ValueError, match="captured plan"): + reader.read_batch(plan.splits()) + assert [result.positions for result in reader.read_batch([])] == [()] * len(QUERIES) diff --git a/paimon-python/pypaimon/tests/ray_search_execution_test.py b/paimon-python/pypaimon/tests/ray_search_execution_test.py index ec337c55067a..2fef666cfb69 100644 --- a/paimon-python/pypaimon/tests/ray_search_execution_test.py +++ b/paimon-python/pypaimon/tests/ray_search_execution_test.py @@ -205,17 +205,28 @@ def run(): cancel.assert_called_once_with('pending') -@pytest.mark.parametrize('family', ['vector', 'full-text']) -def test_python_primary_key_readers_reject_unsupported_executor_explicitly(family): +def test_python_primary_key_full_text_rejects_unsupported_executor_explicitly(): from pypaimon.execution.python.read.search.primary_key_full_text_read import PrimaryKeyFullTextRead - from pypaimon.execution.python.read.search.primary_key_vector_read import PrimaryKeyVectorRead - reader_type = PrimaryKeyVectorRead if family == 'vector' else PrimaryKeyFullTextRead - reader = object.__new__(reader_type) - with pytest.raises(ValueError, match='Ray %s search supports only data-evolution tables' % family): + reader = object.__new__(PrimaryKeyFullTextRead) + with pytest.raises(ValueError, match='Ray full-text search supports only data-evolution tables'): reader.read_plan(object(), executor=RaySearchTaskExecutor(2)) +def test_python_primary_key_vector_read_accepts_executor(): + from pypaimon.common.options.core_options import GlobalIndexSearchMode + from pypaimon.execution.python.read.search.primary_key_vector_read import PrimaryKeyVectorRead + from pypaimon.table.source.search_types import PrimaryKeyVectorScanPlan + + table = SimpleNamespace(options=SimpleNamespace( + primary_key_vector_index_type=lambda column: 'ivf-flat', + vector_index_search_mode=lambda: GlobalIndexSearchMode.FAST)) + reader = PrimaryKeyVectorRead(table, 1, SimpleNamespace(name='embedding'), [1.]) + result = reader.read_plan(PrimaryKeyVectorScanPlan(0, []), executor=RaySearchTaskExecutor(2)) + assert result.snapshot_id == 0 + assert not result.positions + + class _CallbackDriver: def __init__(self): self.items = [] diff --git a/paimon-python/pypaimon/tests/vector_search_factory_selection_test.py b/paimon-python/pypaimon/tests/vector_search_factory_selection_test.py new file mode 100644 index 000000000000..5c691dc16a57 --- /dev/null +++ b/paimon-python/pypaimon/tests/vector_search_factory_selection_test.py @@ -0,0 +1,135 @@ +# 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 pyarrow as pa +import pytest + +from pypaimon import CatalogFactory, Schema +from pypaimon.execution.python.read.search.primary_key_batch_vector_read import BatchPrimaryKeyVectorRead +from pypaimon.execution.python.read.search.primary_key_vector_read import PrimaryKeyVectorRead +from pypaimon.execution.python.read.search.primary_key_vector_scan import PrimaryKeyVectorScan +from pypaimon.execution.python.read.search.vector_search_read import ( + BatchVectorSearchReadImpl, DataEvolutionVectorRead, +) +from pypaimon.execution.python.read.search.vector_search_scan import DataEvolutionVectorScan + + +ARROW_SCHEMA = pa.schema([ + ("id", pa.int64()), ("pt", pa.int32()), + ("embedding", pa.list_(pa.float32(), 2)), ("other", pa.list_(pa.float32(), 2)), +]) + + +@pytest.fixture(params=[False, True], ids=["data-evolution", "primary-key"]) +def table(tmp_path, request): + options = { + "file.format": "parquet", "read.native.enabled": "false", "write.native.enabled": "false", + "scan.native-plan.enabled": "false", "vector-index.search-mode": "full", + } + if request.param: + options.update({"bucket": "2", "pk-vector.index.columns": "embedding", + "fields.embedding.pk-vector.index.type": "ivf-flat"}) + else: + options.update({"row-tracking.enabled": "true", "data-evolution.enabled": "true"}) + catalog = CatalogFactory.create({"warehouse": str(tmp_path)}) + catalog.create_database("default", True) + catalog.create_table("default.vectors", Schema.from_pyarrow_schema( + ARROW_SCHEMA, partition_keys=["pt"], primary_keys=["id", "pt"] if request.param else [], + options=options), False) + return catalog.get_table("default.vectors") + + +def _builders(table, column="embedding"): + single = (table.new_vector_search_builder().with_vector_column(column) + .with_query_vector([1., 0.]).with_limit(2)) + batch = (table.new_batch_vector_search_builder().with_vector_column(column) + .with_query_vectors([[1., 0.], [0., 1.]]).with_limit(2)) + return single, batch + + +@pytest.mark.parametrize("column", ["embedding", "other"]) +def test_public_builders_select_matching_scan_and_read_types(table, column): + single, batch = _builders(table, column) + uses_pk_index = table.is_primary_key_table and column == "embedding" + scan_type, read_type, batch_type = ( + (PrimaryKeyVectorScan, PrimaryKeyVectorRead, BatchPrimaryKeyVectorRead) if uses_pk_index else + (DataEvolutionVectorScan, DataEvolutionVectorRead, BatchVectorSearchReadImpl)) + assert type(single.new_vector_search_scan()) is scan_type + assert type(batch.new_vector_search_scan()) is scan_type + assert type(single.new_vector_search_read()) is read_type + assert type(batch.new_batch_vector_search_read()) is batch_type + + +def test_public_builder_parameters_are_captured_independently(table): + single, batch = _builders(table) + queries = [[1., 0.], [0., 1.]] + single.with_query_vector(queries[0]) + batch.with_query_vectors(queries) + predicates = table.new_read_builder().new_predicate_builder() + row_filter = predicates.greater_or_equal("id", 2) + scans, readers = [], [] + for builder in (single, batch): + builder.with_filter(row_filter).with_partition_filter(predicates.equal("pt", 1)) + builder.with_options({"ivf.nprobe": "2", "metric": "l2"}) + scans.append(builder.new_vector_search_scan()) + readers.append(single.new_vector_search_read()) + readers.append(batch.new_batch_vector_search_read()) + for builder in (single, batch): + builder.with_limit(99).with_vector_column("other").with_option("ivf.nprobe", "8") + builder.with_filter(predicates.greater_or_equal("id", 3)) + queries[0][0] = 99. + queries.append([9., 9.]) + + for execution in scans + readers: + assert execution._table is table + assert execution._vector_column.name == "embedding" + assert execution._filter == row_filter + assert execution._partition_filter.field == "pt" + assert execution._partition_filter.index == 0 + assert execution._options == {"ivf.nprobe": "2", "metric": "l2"} + assert [reader._limit for reader in readers] == [2, 2] + assert readers[0]._query_vector == [1., 0.] + assert readers[1]._query_vectors == [[1., 0.], [0., 1.]] + + +def _append(table, row_id): + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.write_arrow(pa.Table.from_pylist([ + {"id": row_id, "pt": 1, "embedding": [float(row_id), 0.], "other": [0., 1.]}, + ], schema=ARROW_SCHEMA)) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + + +def test_public_builders_preserve_the_selected_snapshot(table): + _append(table, 1) + snapshot = table.snapshot_manager().get_latest_snapshot() + historical = table.copy({"scan.snapshot-id": str(snapshot.id)}) + single, batch = _builders(historical) + scans = [single.new_vector_search_scan(), batch.new_vector_search_scan()] + readers = [single.new_vector_search_read(), batch.new_batch_vector_search_read()] + _append(table, 2) + + assert table.snapshot_manager().get_latest_snapshot().id > snapshot.id + for scan, reader in zip(scans, readers): + assert scan.scan().snapshot().id == snapshot.id + assert reader._table is historical + assert reader._table.options.options.to_map()["scan.snapshot-id"] == str(snapshot.id)