From 53efc5087251fe4e8d07b0de92da53050a2158e4 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Wed, 21 Oct 2026 20:23:40 +0800 Subject: [PATCH 1/9] [Python] Route distributed search through backend task executors --- docs/docs/pypaimon/multimodal-search.md | 43 ++- paimon-python/dev/test_packaging.py | 6 + paimon-python/pypaimon/execution/README.md | 10 + .../native/read/native_hybrid_search.py | 16 +- .../native/read/native_search_plan.py | 36 +- .../execution/native/read/search_execution.py | 107 ++++++ .../search/distributed_full_text_search.py | 115 ++++++ .../python/read/search/distributed_search.py | 47 +++ .../read/search/distributed_vector_search.py | 281 ++++++++++++++ .../python/read/search/full_text_read.py | 8 +- .../python/read/search/hybrid_search.py | 46 ++- .../read/search/primary_key_full_text_read.py | 4 +- .../read/search/primary_key_vector_read.py | 4 +- .../python/read/search/vector_search_read.py | 4 + paimon-python/pypaimon/multimodal/query.py | 8 +- .../pypaimon/ray/batch_vector_search.py | 137 +------ .../pypaimon/ray/full_text_search.py | 150 +------- .../pypaimon/ray/search_execution.py | 201 ++++++++++ .../pypaimon/ray/vector_index_build.py | 4 +- paimon-python/pypaimon/ray/vector_search.py | 269 +------------- .../table/source/hybrid_search_builder.py | 13 + .../pypaimon/table/source/search_dispatch.py | 27 +- .../pypaimon/table/source/search_execution.py | 73 ++++ .../pypaimon/table/source/search_types.py | 34 +- .../pypaimon/tests/native_search_plan_test.py | 60 ++- .../tests/ray_batch_vector_refinement_test.py | 27 +- .../tests/ray_batch_vector_search_test.py | 48 +-- .../tests/ray_full_text_search_test.py | 22 +- .../tests/ray_hybrid_parallel_test.py | 71 ++-- .../tests/ray_search_execution_test.py | 345 ++++++++++++++++++ .../tests/ray_vector_filter_exactness_test.py | 34 +- .../tests/ray_vector_refinement_test.py | 29 +- .../pypaimon/tests/ray_vector_search_test.py | 68 ++-- .../tests/search_execution_test_helpers.py | 33 ++ .../tests/search_task_execution_test.py | 219 +++++++++++ 35 files changed, 1850 insertions(+), 749 deletions(-) create mode 100644 paimon-python/pypaimon/execution/native/read/search_execution.py create mode 100644 paimon-python/pypaimon/execution/python/read/search/distributed_full_text_search.py create mode 100644 paimon-python/pypaimon/execution/python/read/search/distributed_search.py create mode 100644 paimon-python/pypaimon/execution/python/read/search/distributed_vector_search.py create mode 100644 paimon-python/pypaimon/ray/search_execution.py create mode 100644 paimon-python/pypaimon/table/source/search_execution.py create mode 100644 paimon-python/pypaimon/tests/ray_search_execution_test.py create mode 100644 paimon-python/pypaimon/tests/search_execution_test_helpers.py create mode 100644 paimon-python/pypaimon/tests/search_task_execution_test.py diff --git a/docs/docs/pypaimon/multimodal-search.md b/docs/docs/pypaimon/multimodal-search.md index 4470bf3d4e99..9daa8153148f 100644 --- a/docs/docs/pypaimon/multimodal-search.md +++ b/docs/docs/pypaimon/multimodal-search.md @@ -200,8 +200,8 @@ data. Historical snapshot and deletion semantics remain the same. Plain 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 and hybrid queries are not supported -by this execution mode. +data-evolution tables; primary-key tables are not supported by this execution +mode. Full-text and hybrid queries are described below. Install the same PyPaimon and index dependencies on the driver and workers: @@ -235,6 +235,12 @@ the number of in-flight Ray tasks, with a default of 4. `ray_remote_args` suppli Ray task options, including resources and retry settings; `num_returns` is managed by PyPaimon. These two arguments require `execution="ray"`. +Ray schedules tasks from the selected search backend. With Native reads enabled +on a supported REST table, Rust performs the query, filtering and ranking on the +driver and workers. Install the same Native SDK on every node. Selecting Ray +does not force Python search execution. A failure after execution starts is +propagated without rerunning the query through another backend. + For single and batch vector queries, `to_pandas()` and `to_list()` accept the same `execution`, `concurrency`, and `ray_remote_args` arguments as `to_arrow()`. For example, use `query.to_pandas(execution="ray", concurrency=4)` to return a @@ -274,7 +280,7 @@ 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, workers evaluate scalar predicates and build +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 share one Ray scalar-index evaluation; its result is passed directly to dependent tasks without materializing the bitmap on the driver. Candidate verification @@ -282,6 +288,9 @@ reads only filter columns and row IDs in the task's shard range, without scannin the scalar indexes again. The final shard filter is reused across a batch's query vectors. This lets expensive verification scale with search concurrency; partially overlapping scalar index inputs can still cause repeated index reads. +The Native backend resolves shared snapshot filters in Rust before distributing +index and data-reading tasks. Query matrices and other shared task context are +transported once per task group. Index search, raw scans, refinement, lookup, and task retries use the same read snapshot. The existing `snapshot_id` and `tag_name` arguments to `search()` also @@ -324,11 +333,35 @@ uncovered corpus, preserving existing BM25 statistics; that corpus must fit the worker's memory. Pre-filters and deletion vectors keep the local ranking semantics. Hybrid routes share one resolved snapshot, use Ray for their shard work, and -reuse the existing RRF, MRR, or weighted-score fusion on the driver. Routes run -successively so `concurrency` bounds the whole query. Final Arrow/Pandas row lookup +reuse the existing RRF, MRR, or weighted-score fusion on the driver. Concurrent +routes share the query's task budget, so `concurrency` bounds the whole query. +Final Arrow/Pandas row lookup remains on the driver. Native resources close on workers; task errors propagate and pending read tasks are cancelled. +### Scheduling a search plan + +Lower-level vector, batch-vector and full-text readers accept an optional +`SearchTaskExecutor` in `read_plan()` or `read_batch_plan()`. The reader keeps +ownership of the complete snapshot-scoped query, including global candidate +selection and refinement. A scheduler only places the backend's serializable +tasks and returns their results: + +```python +from pypaimon.ray.search_execution import RaySearchTaskExecutor + +builder = (table.new_batch_vector_search_builder() + .with_vector_column("embedding") + .with_query_vectors([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) + .with_limit(10)) +plan = builder.new_vector_search_scan().scan() +results = builder.new_batch_vector_search_read().read_batch_plan( + plan, executor=RaySearchTaskExecutor(concurrency=4)) +``` + +Hybrid builders expose `execute(executor=...)` for the same scheduler. +The existing `execute_local()` methods retain their local defaults. + ## Search Hybrid Use `search_hybrid` to combine vector and full-text routes, then rerank the diff --git a/paimon-python/dev/test_packaging.py b/paimon-python/dev/test_packaging.py index 8803119e9166..89fabc90039d 100644 --- a/paimon-python/dev/test_packaging.py +++ b/paimon-python/dev/test_packaging.py @@ -122,6 +122,12 @@ def test_distribution_contents(self): for name in ("pypaimon/sample/data/data.jsonl", "pypaimon/_version.py"): self.assertIn(name, names) execution_files = ( + "pypaimon/table/source/search_execution.py", + "pypaimon/ray/search_execution.py", + "pypaimon/execution/native/read/search_execution.py", + "pypaimon/execution/python/read/search/distributed_search.py", + "pypaimon/execution/python/read/search/distributed_vector_search.py", + "pypaimon/execution/python/read/search/distributed_full_text_search.py", "pypaimon/execution/python/map_blob_key_serializer.py", "pypaimon/execution/python/write/external_path_provider.py", "pypaimon/execution/python/read/predicate_json_parser.py", diff --git a/paimon-python/pypaimon/execution/README.md b/paimon-python/pypaimon/execution/README.md index 48c0d05ea6b3..532eb5c5038e 100644 --- a/paimon-python/pypaimon/execution/README.md +++ b/paimon-python/pypaimon/execution/README.md @@ -144,6 +144,16 @@ Ray package. Public Ray helper entry points delegate to those same helpers. System-table output applies common Arrow filtering/projection to `SystemSplit` results from either metadata provider without creating a Python file reader. +Search readers accept an optional `SearchTaskExecutor` when executing a complete +plan. `SearchTask` and its dependency/context transport contract stay in +`table/source/search_execution`; Ray only implements bounded scheduling, +context sharing and cancellation. It does not inherit a Python search reader +or select Python on behalf of a Native reader. Python's distributed adapters +own their index/filter/raw/refinement operations. Native adapters transport +opaque shared context and task bytes; Core owns both their interpretation and +global query coordination. The global candidate limit is applied before +refinement, and full-text raw fallback retains one corpus for BM25 statistics. + ## Corresponding execution APIs Compare the operation's inputs, results and lifecycle when changing a backend. diff --git a/paimon-python/pypaimon/execution/native/read/native_hybrid_search.py b/paimon-python/pypaimon/execution/native/read/native_hybrid_search.py index b40a7676d211..d7c485fcfac3 100644 --- a/paimon-python/pypaimon/execution/native/read/native_hybrid_search.py +++ b/paimon-python/pypaimon/execution/native/read/native_hybrid_search.py @@ -27,7 +27,7 @@ from pypaimon.execution.native.write.native_commit import _rest_catalog_supported -def try_native_hybrid_search(builder): +def try_native_hybrid_search(builder, executor=None): """Qualify before execution; core owns one snapshot, routes and fusion.""" table = builder._table if type(table) is not FileStoreTable: @@ -41,6 +41,8 @@ def try_native_hybrid_search(builder): or not _rest_catalog_supported(table)): return None builder._validate_search() + if executor is not None and not table.options.data_evolution_enabled(): + raise ValueError('Distributed search supports only data-evolution tables') if any(route.is_full_text() for route in builder._routes): # Preserve the same typed analyzer normalization used by native full-text. options = {'full-text.' + key: value for key, value in NativeFullTextIndexOptions.from_options( @@ -63,4 +65,14 @@ def try_native_hybrid_search(builder): # Qualify strict literal support before executing any route or ranking. return None # Propagate execution failures: retrying Python can select a newer snapshot. - return DictBasedScoredIndexResult(native.execute_local().row_ids()) + if executor is None: + result = native.execute_local() + else: + from pypaimon.execution.native.read.search_execution import native_task_executor + callback = native_task_executor(table, executor) + try: + result = native.execute(executor=callback) + finally: + callback.cancel() + callback.raise_if_failed() + return DictBasedScoredIndexResult(result.row_ids()) diff --git a/paimon-python/pypaimon/execution/native/read/native_search_plan.py b/paimon-python/pypaimon/execution/native/read/native_search_plan.py index e55f812243ba..be49b2660197 100644 --- a/paimon-python/pypaimon/execution/native/read/native_search_plan.py +++ b/paimon-python/pypaimon/execution/native/read/native_search_plan.py @@ -163,12 +163,38 @@ def __getstate__(self): state['_native'] = None return state - def read_plan(self, plan): - return _result(self._table, self._native.read_plan(plan), self._primary_key) - - def read_batch_plan(self, plan): + def read_plan(self, plan, executor=None): + if executor is None: + result = self._native.read_plan(plan) + else: + self._validate_task_execution() + from pypaimon.execution.native.read.search_execution import native_task_executor + callback = native_task_executor(self._table, executor) + try: + result = self._native.read_plan(plan, executor=callback) + finally: + callback.cancel() + callback.raise_if_failed() + return _result(self._table, result, self._primary_key) + + def read_batch_plan(self, plan, executor=None): + if executor is None: + results = self._native.read_batch_plan(plan) + else: + self._validate_task_execution() + from pypaimon.execution.native.read.search_execution import native_task_executor + callback = native_task_executor(self._table, executor) + try: + results = self._native.read_batch_plan(plan, executor=callback) + finally: + callback.cancel() + callback.raise_if_failed() return [_result(self._table, result, False) - for result in self._native.read_batch_plan(plan)] + for result in results] + + def _validate_task_execution(self): + if self._primary_key or not self._table.options.data_evolution_enabled(): + raise ValueError('Distributed search supports only data-evolution tables') def resolve_plan(self, plan): # Search modes choose source work at scan time. A new reader may use diff --git a/paimon-python/pypaimon/execution/native/read/search_execution.py b/paimon-python/pypaimon/execution/native/read/search_execution.py new file mode 100644 index 000000000000..13241198d1d4 --- /dev/null +++ b/paimon-python/pypaimon/execution/native/read/search_execution.py @@ -0,0 +1,107 @@ +# 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. + +"""Adapt opaque Rust search tasks to the shared scheduling contract.""" + +from concurrent.futures import CancelledError +from threading import Event, Lock + +from pypaimon.execution.native.table import _native_table +from pypaimon.table.source.search_execution import SearchTask + + +def _execute_task(context, payload): + # Core restores the captured schema/snapshot and owns file readers. Only + # ordinary table configuration and opaque bytes cross the worker boundary. + table, shared = context + return _native_table(table)._execute_search_task(payload, shared) + + +def native_task_executor(table, executor): + return _NativeTaskExecutor(table, executor) + + +class _NativeTaskExecutor: + """One cancellation scope shared by all Core callbacks for a query.""" + + def __init__(self, table, executor): + self._table = table + self._executor = executor.with_concurrency(executor.concurrency) + self._cancelled = Event() + self._failure_lock = Lock() + self._failure = None + + def cancel(self): + with self._failure_lock: + self._cancelled.set() + + def _fail(self, failure): + # Save the cause before waking sibling callbacks: their synthetic + # cancellation may reach the Rust boundary before this callback closes. + with self._failure_lock: + if (self._failure is None and not self._cancelled.is_set() + and not isinstance(failure, GeneratorExit)): + self._failure = failure + # Once Core or another callback cancels the query, late failures + # and cleanup errors must not replace its existing primary error. + self._cancelled.set() + + def raise_if_failed(self): + with self._failure_lock: + failure = self._failure + if failure is not None: + raise failure + + def __call__(self, shared, payloads): + try: + if self._cancelled.is_set(): + raise CancelledError('Search query cancelled') + with self._executor.cancellation_scope(self._cancelled): + yield from self._execute(shared, payloads) + except BaseException as failure: + self._fail(failure) + raise + + def _execute(self, shared, payloads): + context = (self._table, shared) + tasks = [SearchTask(_execute_task, context, payload) for payload in payloads] + results = self._executor.map(tasks, ordered=True) + completed = 0 + failed = False + try: + for ordinal, value in results: + if self._cancelled.is_set(): + raise CancelledError('Search query cancelled') + if ordinal != completed or completed >= len(tasks): + raise ValueError('Search task executor did not preserve task order') + completed += 1 + yield value + if completed != len(tasks): + raise ValueError('Search task executor returned an incomplete result sequence') + except BaseException as failure: + failed = True + self._fail(failure) + raise + finally: + close = getattr(results, 'close', None) + if close is not None: + try: + close() + except BaseException as failure: + self._fail(failure) + if not failed: + raise diff --git a/paimon-python/pypaimon/execution/python/read/search/distributed_full_text_search.py b/paimon-python/pypaimon/execution/python/read/search/distributed_full_text_search.py new file mode 100644 index 000000000000..fa1497e60cc8 --- /dev/null +++ b/paimon-python/pypaimon/execution/python/read/search/distributed_full_text_search.py @@ -0,0 +1,115 @@ +# 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. + +"""Python full-text tasks preserve corpus statistics and plan-ordered merging.""" + +from contextlib import closing + +from pypaimon.globalindex.global_index_result import GlobalIndexResult +from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult +from pypaimon.execution.python.read.search.full_text_read import ( + DataEvolutionFullTextRead, _index_candidate_limit, _index_type, _raw_row_ranges, _split_search_splits, +) +from pypaimon.execution.python.read.search import global_index_live_row_filter +from pypaimon.execution.python.read.search.distributed_vector_search import _index_tasks, _scores +from pypaimon.table.source.search_execution import SearchTask +from pypaimon.utils.range import Range +from pypaimon.utils.roaring_bitmap import RoaringBitmap64 + + +class DistributedFullTextRead(DataEvolutionFullTextRead): + def __init__(self, reader, executor): + super().__init__(reader._table, reader._limit, reader._text_columns, reader._query, + reader._partition_filter, reader._filter) + self._executor = executor + + def _worker_reader(self): + reader = DataEvolutionFullTextRead( + self._table.copy({"global-index.enabled": "false"}), self._limit, self._text_columns, self._query, + self._partition_filter, self._filter) + reader._snapshot = self._snapshot + return reader + + def read(self, splits): + # Keep bitmap construction and candidate data reads off the driver. + index_splits, raw_splits = _split_search_splits(splits) + indexed = self._eval_column_query({self._text_columns[0].name: index_splits}, None) + raw = self._read_raw_search(_raw_row_ranges(raw_splits), _index_type(index_splits)) + return indexed.or_(raw).top_k(self._limit) + + def _eval_column_query(self, splits_by_column, live_rows): + items = splits_by_column.get(self._text_columns[0].name, []) + if not items: + return DictBasedScoredIndexResult({}) + merged = {} + reader = self._worker_reader() + reader._limit = _index_candidate_limit(items, self._limit) + work = _index_tasks(_search_index, reader, items, self._filter, _prepare_scalar_filter) + with closing(self._executor.map(work, ordered=True)) as tasks: + for _, scores in tasks: + # Duplicate row IDs keep the first planned shard's score. + for row_id, score in scores.items(): + merged.setdefault(row_id, score) + return DictBasedScoredIndexResult(merged).top_k(self._limit) + + def _read_raw_search(self, ranges, index_type, include_row_ids=None): + if not ranges: + return DictBasedScoredIndexResult({}) + include = None if include_row_ids is None else include_row_ids.serialize() + # Splitting this corpus would change document frequencies and BM25. + # Keep the existing local corpus intact, but read/build/search on a worker. + work = [SearchTask(_search_raw, self._worker_reader(), (ranges, index_type, include))] + with closing(self._executor.map(work)) as tasks: + for _, scores in tasks: + return DictBasedScoredIndexResult(scores) + + +def _search_index(reader, split, scalar_filter=None): + live = global_index_live_row_filter.live_rows( + reader._table, reader._partition_filter, reader._snapshot, + row_ranges=[Range(split.row_range_start, split.row_range_end)]) + if scalar_filter is None: + include = reader._matched_indexed_rows([split], live) + else: + bitmap, exact = scalar_filter + result = None if bitmap is None else GlobalIndexResult.create( + RoaringBitmap64.deserialize(bitmap), is_exact=exact) + include = reader._matched_indexed_rows([split], live, result) + if include is not None and include.is_empty(): + return {} + return _scores(reader._eval(split.row_range_start, split.row_range_end, + split.full_text_index_files, include).result()) + + +def _prepare_scalar_filter(reader, split): + result = reader._scalar_index_result(split.scalar_index_files, reader._snapshot) + return (None, False) if result is None else (result.results().serialize(), result.is_exact()) + + +def _search_raw(reader, item): + ranges, index_type, include = item + include = None if include is None else RoaringBitmap64.deserialize(include) + if reader._filter is not None: + live = global_index_live_row_filter.live_rows( + reader._table, reader._partition_filter, reader._snapshot, row_ranges=ranges) + candidates = GlobalIndexResult.from_ranges(ranges).results() + if live is not None: + candidates = RoaringBitmap64.and_(candidates, live) + if include is not None: + candidates = RoaringBitmap64.and_(candidates, include) + include = reader._matching_rows(candidates) + return _scores(reader._read_raw_search(ranges, index_type, include)) diff --git a/paimon-python/pypaimon/execution/python/read/search/distributed_search.py b/paimon-python/pypaimon/execution/python/read/search/distributed_search.py new file mode 100644 index 000000000000..d8af398a8c68 --- /dev/null +++ b/paimon-python/pypaimon/execution/python/read/search/distributed_search.py @@ -0,0 +1,47 @@ +# 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. + +"""Attach distributed task execution to Python search readers.""" + +import math + + +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): + if (type(reader) is not DataEvolutionFullTextRead + or not reader._table.options.data_evolution_enabled()): + 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) diff --git a/paimon-python/pypaimon/execution/python/read/search/distributed_vector_search.py b/paimon-python/pypaimon/execution/python/read/search/distributed_vector_search.py new file mode 100644 index 000000000000..95618ee8eb7f --- /dev/null +++ b/paimon-python/pypaimon/execution/python/read/search/distributed_vector_search.py @@ -0,0 +1,281 @@ +# 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. + +"""Python search tasks and global result coordination for distributed executors.""" + +from collections import Counter +from contextlib import closing +import math + +from pypaimon.globalindex.batch_vector_search import BatchVectorSearch +from pypaimon.globalindex.global_index_result import GlobalIndexResult +from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult +from pypaimon.execution.python.read.search.vector_search_read import ( + BatchVectorSearchReadImpl, DataEvolutionVectorRead, + _filtered_raw_row_ranges, _offer_score, _scored_result, _split_search_splits, +) +from pypaimon.execution.python.read.search import global_index_live_row_filter +from pypaimon.table.source.search_execution import SearchTask +from pypaimon.utils.range import Range +from pypaimon.utils.roaring_bitmap import RoaringBitmap64 + + +class DistributedVectorSearchRead(DataEvolutionVectorRead): + """Replace split execution while retaining local merge/refinement semantics.""" + + 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 _prepare_search_splits(self, splits, snapshot): + # Scalar evaluation and candidate verification belong to the workers. + return _split_search_splits(splits) + + def _search_index_splits(self, splits, query, search_limit, snapshot, batch=False): + if not splits: + return + # Send predicates and metadata, leaving each worker to build and consume + # its own filter. Batch queries reuse the same filter for every vector. + context = (self._table, self._vector_column, query, search_limit, self._options, + self._filter, self._partition_filter, snapshot) + work = _index_tasks(_search_index_split, context, splits, self._filter) + with closing(self._executor.map(work, ordered=True)) as tasks: + for _, (metric, scores) in tasks: + if metric is not None: + self._set_index_metric(metric) + # Keep plan order, including duplicate-row precedence, regardless + # of worker completion order. Do not refine or truncate per worker. + yield DictBasedScoredIndexResult(scores) + del scores + + def _read_raw_search(self, raw_row_ranges, pre_filter, query_vector, + index_type=None, include_filter=True, + score_candidates=None, snapshot=None): + ranges = _filtered_raw_row_ranges(raw_row_ranges, pre_filter) + if score_candidates is not None: + # Refine only the globally selected candidates. Intersect before + # planning so workers neither read nor score other rows. + ranges = _filtered_raw_row_ranges(ranges, score_candidates.to_range_list()) + if not ranges: + return DictBasedScoredIndexResult({}) + table_read, splits = self._plan_raw_read(ranges, include_filter, snapshot) + context = (table_read, self._vector_column, query_vector, + self._limit, self._search_metric(index_type)) + heap = [] + # Keep each planned data split intact: its files may be vertical column + # groups that must be read together to reconstruct a logical row. + work = [SearchTask(_search_raw_split, context, split) for split in splits] + with closing(self._executor.map(work)) as tasks: + for _, scores in tasks: + for row_id, score in scores.items(): + _offer_score(heap, self._limit, row_id, score) + return _scored_result(heap) + + +class DistributedBatchVectorSearchRead(BatchVectorSearchReadImpl): + + 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 _prepare_search_splits(self, splits, snapshot): + # Scalar evaluation and candidate verification belong to the workers. + return _split_search_splits(splits) + + def _search_index_splits(self, splits, query, search_limit, snapshot, batch=False): + if not splits: + return + # Send predicates and metadata, leaving each worker to build and consume + # its own filter. Batch queries reuse the same filter for every vector. + context = (self._table, self._vector_column, query, search_limit, self._options, + self._filter, self._partition_filter, snapshot) + work = _index_tasks(_search_batch_index_split, context, splits, self._filter) + with closing(self._executor.map(work, ordered=True)) as tasks: + for _, (metric, scores) in tasks: + if metric is not None: + self._set_index_metric(metric) + # Retain plan order for duplicate IDs and global per-query selection. + yield [DictBasedScoredIndexResult(values) for values in scores] + del scores + + def _read_raw_batch_search(self, raw_row_ranges, pre_filter, index_type=None, snapshot=None): + heaps = [[] for _ in self._query_vectors] + ranges = _filtered_raw_row_ranges(raw_row_ranges, pre_filter) + if not ranges: + return [_scored_result(heap) for heap in heaps] + table_read, splits = self._plan_raw_read(ranges, True, snapshot) + context = (table_read, self._vector_column, self._query_vectors, + self._limit, self._search_metric(index_type)) + work = [SearchTask(_search_batch_raw_split, context, split) for split in splits] + with closing(self._executor.map(work)) as tasks: + for _, results in tasks: + for heap, scores in zip(heaps, results): + for row_id, score in scores.items(): + _offer_score(heap, self._limit, row_id, score) + return [_scored_result(heap) for heap in heaps] + + def _stream_rerank_candidates(self, candidates, union_candidates, query_vectors, + index_type, snapshot): + # Candidates have already been selected globally, separately per query. + queries_by_row = {} + for query_index, result in enumerate(candidates): + for row_id in result.results(): + queries_by_row.setdefault(row_id, []).append(query_index) + table_read, splits = self._plan_raw_read( + union_candidates.to_range_list(), include_filter=False, snapshot=snapshot) + context = (table_read, self._vector_column, query_vectors, self._limit, + self._search_metric(index_type), queries_by_row) + heaps = [[] for _ in query_vectors] + work = [SearchTask(_search_batch_refine_split, context, split) for split in splits] + with closing(self._executor.map(work)) as tasks: + for _, results in tasks: + for heap, scores in zip(heaps, results): + for row_id, score in scores.items(): + _offer_score(heap, self._limit, row_id, score) + return [_scored_result(heap) for heap in heaps] + + +def _search_index_split(context, split, scalar_filter=None): + table, column, query, limit, options, filter_, partition_filter, snapshot = context + # Explicit shard indexes are evaluated below. Candidate verification must + # not rediscover and scan table-wide scalar indexes in every worker. + table = table.copy({"global-index.enabled": "false"}) + reader = DataEvolutionVectorRead( + table, limit, column, query, filter_, partition_filter, options) + include = _worker_pre_filter(reader, split, snapshot, scalar_filter) + result = reader._eval( + split.row_range_start, split.row_range_end, + split.vector_index_files, query, limit, include).result() + # The reader closes its native resources before returning a completed result. + # Return metric even for zero hits, so inconsistent indexes cannot be hidden + # by a scalar filter and raw fallback uses the persisted index metric. + return reader._index_metric, _scores(result) + + +def _worker_pre_filter(reader, split, snapshot, scalar_filter): + live_rows = global_index_live_row_filter.live_rows( + reader._table, reader._partition_filter, snapshot, + row_ranges=[Range(split.row_range_start, split.row_range_end)]) + if scalar_filter is None: + matched = reader._scalar_matched_rows([split], snapshot) + else: + bitmap, exact = scalar_filter + result = None if bitmap is None else GlobalIndexResult.create( + RoaringBitmap64.deserialize(bitmap), is_exact=exact) + matched = reader._refine_scalar_result([split], result, snapshot) + filters = reader._combine_pre_filters([split], live_rows, matched) + return filters[0] if filters else None + + +def _prepare_scalar_filter(context, split): + table, column, query, limit, options, filter_, partition_filter, snapshot = context + reader = DataEvolutionVectorRead(table, limit, column, query, filter_, partition_filter, options) + result = reader._scalar_index_result([split], snapshot) + return (None, False) if result is None else (result.results().serialize(), result.is_exact()) + + +def _index_tasks(operation, context, splits, predicate, scalar_operation=None): + """Share identical scalar work while each shard verifies its own candidates.""" + scalar_operation = scalar_operation or _prepare_scalar_filter + + def key(split): + return tuple(sorted(file.file_name for file in split.scalar_index_files)) + + counts = Counter(key(split) for split in splits) if predicate is not None else {} + shared = {} + tasks = [] + for split in splits: + files = key(split) if counts else None + dependencies = () + if files and counts.get(files, 0) > 1: + if files not in shared: + shared[files] = SearchTask(scalar_operation, context, split) + dependencies = (shared[files],) + tasks.append(SearchTask(operation, context, split, dependencies)) + return tasks + + +def _search_raw_split(context, split): + from pypaimon.read.read_support import _ClosableArrowBatchReader + + table_read, column, query, limit, metric = context + scorer = DataEvolutionVectorRead(table_read.table, limit, column, query) + reader, batches = table_read._new_arrow_batch_reader([split]) + with _ClosableArrowBatchReader(reader, batches) as batch_reader: + return _scores(scorer._score_raw_batches(batch_reader, query, metric, reject_nan=True)) + + +def _scores(result): + if result is None: + return {} + getter = result.score_getter() + scores = {row_id: getter(row_id) for row_id in result.results()} + if any(math.isnan(score) for score in scores.values()): + raise ValueError("Ray vector search cannot rank NaN scores.") + return scores + + +def _search_batch_index_split(context, split, scalar_filter=None): + table, column, queries, limit, options, filter_, partition_filter, snapshot = context + # Explicit shard indexes are evaluated below. Candidate verification must + # not rediscover and scan table-wide scalar indexes in every worker. + table = table.copy({"global-index.enabled": "false"}) + scorer = BatchVectorSearchReadImpl( + table, limit, column, queries, filter_, partition_filter, options) + include = _worker_pre_filter(scorer, split, snapshot, scalar_filter) + reader, offset_reader = scorer._open_offset_reader( + split.vector_index_files, split.row_range_start, split.row_range_end) + try: + # Bound native query-matrix scratch to at most 64 queries / 4 MiB + # (or one vector), keeping the shard open across all query blocks. + block_size = max(1, min(64, (1 << 20) // max(1, len(queries[0])))) + results = [] + for start in range(0, len(queries), block_size): + search = BatchVectorSearch(queries[start:start + block_size], limit, column.name, + include_row_ids=include, options=options) + results.extend(_batch_scores(result) for result in offset_reader.visit_batch_vector_search(search).result()) + return scorer._index_metric, results + finally: + reader.close() + + +def _search_batch_raw_split(context, split): + from pypaimon.read.read_support import _ClosableArrowBatchReader + + table_read, column, queries, limit, metric = context + scorer = BatchVectorSearchReadImpl(table_read.table, limit, column, queries) + reader, batches = table_read._new_arrow_batch_reader([split]) + with _ClosableArrowBatchReader(reader, batches) as batch_reader: + return [_batch_scores(result) for result in scorer._score_raw_batch_queries(batch_reader, metric)] + + +def _search_batch_refine_split(context, split): + table_read, column, queries, limit, metric, queries_by_row = context + scorer = BatchVectorSearchReadImpl(table_read.table, limit, column, queries) + # Read the candidate union once, scoring each row only for its own queries. + return [_batch_scores(result) for result in scorer._score_refine_splits( + table_read, [split], queries_by_row, queries, metric)] + + +def _batch_scores(result): + if result is None: + return {} + getter = result.score_getter() + return {row_id: getter(row_id) for row_id in result.results()} diff --git a/paimon-python/pypaimon/execution/python/read/search/full_text_read.py b/paimon-python/pypaimon/execution/python/read/search/full_text_read.py index a7ce89a48481..57bab789d44b 100644 --- a/paimon-python/pypaimon/execution/python/read/search/full_text_read.py +++ b/paimon-python/pypaimon/execution/python/read/search/full_text_read.py @@ -68,8 +68,12 @@ def __init__( self._partition_filter = partition_filter self._filter = filter_ - def read_plan(self, plan: FullTextScanPlan) -> GlobalIndexResult: - reader = copy(self) + def _with_executor(self, executor): + from pypaimon.execution.python.read.search.distributed_search import create_distributed_reader + return create_distributed_reader(self, executor) + + def read_plan(self, plan: FullTextScanPlan, executor=None) -> GlobalIndexResult: + reader = copy(self if executor is None else self._with_executor(executor)) reader._table = global_index_live_row_filter.table_at_snapshot(self._table, plan.snapshot()) reader._snapshot = plan.snapshot() return reader.read(plan.splits()) diff --git a/paimon-python/pypaimon/execution/python/read/search/hybrid_search.py b/paimon-python/pypaimon/execution/python/read/search/hybrid_search.py index 55ead01bd6ae..b53d7cd2314f 100644 --- a/paimon-python/pypaimon/execution/python/read/search/hybrid_search.py +++ b/paimon-python/pypaimon/execution/python/read/search/hybrid_search.py @@ -19,7 +19,8 @@ import heapq from copy import copy -from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait +from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, as_completed, wait +from threading import Event from typing import List from pypaimon.catalog.table_query_auth import reject_search_under_query_auth @@ -35,6 +36,49 @@ _RRF_K = 60.0 +def execute(self, executor=None): + if executor is None: + return execute_local(self) + executor = executor.with_concurrency(executor.concurrency) + # Pin once before route planning moves onto separate driver threads. + routes = self.route_builders() + workers = min(executor.concurrency, len(routes)) + route_concurrency = executor.concurrency // max(1, workers) + + def execute_route(route): + child = executor.with_concurrency(route_concurrency) + builder = route.search_builder + if route.route.is_vector(): + plan = builder.new_vector_search_scan().scan() + return builder.new_vector_search_read().read_plan(plan, executor=child) + plan = builder.new_full_text_scan().scan() + return builder.new_full_text_read().read_plan(plan, executor=child) + + if workers <= 1: + return self.rank([self.to_route_result(route, execute_route(route)) for route in routes]) + + cancelled = Event() + + def run(route): + with executor.cancellation_scope(cancelled): + return execute_route(route) + + results = [None] * len(routes) + with ThreadPoolExecutor(max_workers=workers) as pool: + futures = {} + try: + for ordinal, route in enumerate(routes): + futures[pool.submit(run, route)] = ordinal + for future in as_completed(futures): + ordinal = futures[future] + results[ordinal] = self.to_route_result(routes[ordinal], future.result()) + finally: + cancelled.set() + for future in futures: + future.cancel() + return self.rank(results) + + def execute_local(self) -> ScoredGlobalIndexResult: """Execute hybrid index search locally.""" route_builders = self.route_builders() diff --git a/paimon-python/pypaimon/execution/python/read/search/primary_key_full_text_read.py b/paimon-python/pypaimon/execution/python/read/search/primary_key_full_text_read.py index 5538cd237b14..29f8ee55c932 100644 --- a/paimon-python/pypaimon/execution/python/read/search/primary_key_full_text_read.py +++ b/paimon-python/pypaimon/execution/python/read/search/primary_key_full_text_read.py @@ -49,7 +49,9 @@ def __init__(self, table, limit, text_column, query, definition, self._definition = definition super().__init__(table, limit, text_column, query, partition_filter) - def read_plan(self, plan): + def read_plan(self, plan, executor=None): + if executor is not None: + return self._with_executor(executor).read_plan(plan) if not isinstance(plan, PrimaryKeyFullTextScanPlan): raise ValueError("Primary-key full-text read requires a PrimaryKeyFullTextScanPlan.") futures = [] 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 2da2ae6ea6fb..7bdfbdf84a23 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 @@ -35,7 +35,9 @@ class PrimaryKeyVectorRead(DataEvolutionVectorRead): """Search PK vector payloads and localize payload row ids to data files.""" - def read_plan(self, plan): + def read_plan(self, plan, executor=None): + if executor is not None: + return self._with_executor(executor).read_plan(plan) if not isinstance(plan, PrimaryKeyVectorScanPlan): raise ValueError("Primary-key vector read requires a PrimaryKeyVectorScanPlan.") self._index_metric = None 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 8abc0d6a32fb..56a4587ad5bb 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 @@ -64,6 +64,10 @@ def __init__( self._index_metric = None self._scalar_pre_filter = None + def _with_executor(self, executor): + from pypaimon.execution.python.read.search.distributed_search import create_distributed_reader + return create_distributed_reader(self, executor) + def _search_metric(self, index_type=None): if self._index_metric is not None: return self._index_metric diff --git a/paimon-python/pypaimon/multimodal/query.py b/paimon-python/pypaimon/multimodal/query.py index 8df4287ddc35..9d5375fbcff5 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -607,9 +607,9 @@ def to_ray(self, *, execution="ray", concurrency=None, ray_remote_args=None, until the Dataset has finished executing. """ from pypaimon.ray.search_result import read_search_result - from pypaimon.ray.vector_search import _execution_options + from pypaimon.ray.search_execution import execution_options - concurrency, remote_args = _execution_options(concurrency, ray_remote_args) + concurrency, remote_args = execution_options(concurrency, ray_remote_args) query = self._for_execution() if execution == "ray": result = query._execute_ray(concurrency, remote_args) @@ -861,9 +861,9 @@ def to_ray(self, *, execution="ray", concurrency=None, ray_remote_args=None, """Return one lazy Ray Dataset per query, with distributed row lookup.""" from pypaimon.ray.batch_vector_search import _execute_batch_vector_search from pypaimon.ray.search_result import read_search_result - from pypaimon.ray.vector_search import _execution_options + from pypaimon.ray.search_execution import execution_options - concurrency, remote_args = _execution_options(concurrency, ray_remote_args) + concurrency, remote_args = execution_options(concurrency, ray_remote_args) query = self._for_execution() if execution == "ray": results = _execute_batch_vector_search( diff --git a/paimon-python/pypaimon/ray/batch_vector_search.py b/paimon-python/pypaimon/ray/batch_vector_search.py index 6e2eed1318b7..7a5bdc84f454 100644 --- a/paimon-python/pypaimon/ray/batch_vector_search.py +++ b/paimon-python/pypaimon/ray/batch_vector_search.py @@ -15,140 +15,13 @@ # specific language governing permissions and limitations # under the License. -"""Ray execution of snapshot-pinned batch vector queries.""" +"""Execute batch vector plans through the public search-task protocol.""" -from contextlib import closing - -from pypaimon.globalindex.batch_vector_search import BatchVectorSearch -from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult -from pypaimon.ray.vector_search import ( - _execution_options, _map_tasks, _require_ray, _shared_scalar_filters, _worker_pre_filter, -) -from pypaimon.execution.python.read.search.vector_search_read import ( - BatchVectorSearchReadImpl, _filtered_raw_row_ranges, _offer_score, _scored_result, _split_search_splits, -) +from pypaimon.ray.search_execution import RaySearchTaskExecutor def _execute_batch_vector_search(builder, *, concurrency=None, ray_remote_args=None): - concurrency, remote_args = _execution_options(concurrency, ray_remote_args) - from pypaimon.table.source.search_dispatch import python_search_reader - reader = python_search_reader(builder.new_batch_vector_search_read()) - if (type(reader) is not BatchVectorSearchReadImpl - or not reader._table.options.data_evolution_enabled()): - raise ValueError("Ray vector search supports only data-evolution tables.") - _require_ray() + executor = RaySearchTaskExecutor(concurrency, ray_remote_args) + reader = builder.new_batch_vector_search_read() plan = builder.new_vector_search_scan().scan() - return _RayBatchVectorSearchRead(reader, concurrency, remote_args).read_batch_plan(plan) - - -class _RayBatchVectorSearchRead(BatchVectorSearchReadImpl): - - def __init__(self, reader, concurrency, remote_args): - super().__init__(reader._table, reader._limit, reader._vector_column, reader._query_vectors, - reader._filter, reader._partition_filter, reader._options) - self._concurrency = concurrency - self._remote_args = remote_args - - def _prepare_search_splits(self, splits, snapshot): - # Scalar evaluation and candidate verification belong to the workers. - return _split_search_splits(splits) - - def _search_index_splits(self, splits, query, search_limit, snapshot, batch=False): - # Send predicates and metadata, leaving each worker to build and consume - # its own filter. Batch queries reuse the same filter for every vector. - context = (self._table, self._vector_column, query, search_limit, self._options, - self._filter, self._partition_filter, snapshot) - with _shared_scalar_filters(splits, self._filter, self._remote_args) as prepare, closing(_map_tasks( - _search_batch_index_split, context, splits, self._concurrency, - self._remote_args, True, prepare)) as tasks: - for _, (metric, scores) in tasks: - if metric is not None: - self._set_index_metric(metric) - # Retain plan order for duplicate IDs and global per-query selection. - yield [DictBasedScoredIndexResult(values) for values in scores] - del scores - - def _read_raw_batch_search(self, raw_row_ranges, pre_filter, index_type=None, snapshot=None): - heaps = [[] for _ in self._query_vectors] - ranges = _filtered_raw_row_ranges(raw_row_ranges, pre_filter) - if not ranges: - return [_scored_result(heap) for heap in heaps] - table_read, splits = self._plan_raw_read(ranges, True, snapshot) - context = (table_read, self._vector_column, self._query_vectors, - self._limit, self._search_metric(index_type)) - with closing(_map_tasks( - _search_batch_raw_split, context, splits, self._concurrency, self._remote_args)) as tasks: - for _, results in tasks: - for heap, scores in zip(heaps, results): - for row_id, score in scores.items(): - _offer_score(heap, self._limit, row_id, score) - return [_scored_result(heap) for heap in heaps] - - def _stream_rerank_candidates(self, candidates, union_candidates, query_vectors, - index_type, snapshot): - # Candidates have already been selected globally, separately per query. - queries_by_row = {} - for query_index, result in enumerate(candidates): - for row_id in result.results(): - queries_by_row.setdefault(row_id, []).append(query_index) - table_read, splits = self._plan_raw_read( - union_candidates.to_range_list(), include_filter=False, snapshot=snapshot) - context = (table_read, self._vector_column, query_vectors, self._limit, - self._search_metric(index_type), queries_by_row) - heaps = [[] for _ in query_vectors] - with closing(_map_tasks( - _search_batch_refine_split, context, splits, self._concurrency, self._remote_args)) as tasks: - for _, results in tasks: - for heap, scores in zip(heaps, results): - for row_id, score in scores.items(): - _offer_score(heap, self._limit, row_id, score) - return [_scored_result(heap) for heap in heaps] - - -def _search_batch_index_split(context, split, scalar_filter=None): - table, column, queries, limit, options, filter_, partition_filter, snapshot = context - # Explicit shard indexes are evaluated below. Candidate verification must - # not rediscover and scan table-wide scalar indexes in every worker. - table = table.copy({"global-index.enabled": "false"}) - scorer = BatchVectorSearchReadImpl( - table, limit, column, queries, filter_, partition_filter, options) - include = _worker_pre_filter(scorer, split, snapshot, scalar_filter) - reader, offset_reader = scorer._open_offset_reader( - split.vector_index_files, split.row_range_start, split.row_range_end) - try: - # Bound native query-matrix scratch to at most 64 queries / 4 MiB - # (or one vector), keeping the shard open across all query blocks. - block_size = max(1, min(64, (1 << 20) // max(1, len(queries[0])))) - results = [] - for start in range(0, len(queries), block_size): - search = BatchVectorSearch(queries[start:start + block_size], limit, column.name, - include_row_ids=include, options=options) - results.extend(_scores(result) for result in offset_reader.visit_batch_vector_search(search).result()) - return scorer._index_metric, results - finally: - reader.close() - - -def _search_batch_raw_split(context, split): - from pypaimon.read.read_support import _ClosableArrowBatchReader - - table_read, column, queries, limit, metric = context - scorer = BatchVectorSearchReadImpl(table_read.table, limit, column, queries) - reader, batches = table_read._new_arrow_batch_reader([split]) - with _ClosableArrowBatchReader(reader, batches) as batch_reader: - return [_scores(result) for result in scorer._score_raw_batch_queries(batch_reader, metric)] - - -def _search_batch_refine_split(context, split): - table_read, column, queries, limit, metric, queries_by_row = context - scorer = BatchVectorSearchReadImpl(table_read.table, limit, column, queries) - # Read the candidate union once, scoring each row only for its own queries. - return [_scores(result) for result in scorer._score_refine_splits( - table_read, [split], queries_by_row, queries, metric)] - - -def _scores(result): - if result is None: - return {} - getter = result.score_getter() - return {row_id: getter(row_id) for row_id in result.results()} + return reader.read_batch_plan(plan, executor=executor) diff --git a/paimon-python/pypaimon/ray/full_text_search.py b/paimon-python/pypaimon/ray/full_text_search.py index d3eb37142e94..82c5a9223694 100644 --- a/paimon-python/pypaimon/ray/full_text_search.py +++ b/paimon-python/pypaimon/ray/full_text_search.py @@ -15,155 +15,17 @@ # specific language governing permissions and limitations # under the License. -"""Ray full-text shard execution, preserving local corpus and merge semantics.""" +"""Execute full-text and hybrid searches through public backend interfaces.""" -from concurrent.futures import ThreadPoolExecutor, as_completed -from contextlib import closing -from threading import Event - -from pypaimon.globalindex.global_index_result import GlobalIndexResult -from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult -from pypaimon.ray.vector_search import _execution_options, _map_tasks, _require_ray, _scores, _shared_scalar_filters -from pypaimon.execution.python.read.search.full_text_read import ( - DataEvolutionFullTextRead, _index_candidate_limit, _index_type, _raw_row_ranges, _split_search_splits, -) -from pypaimon.execution.python.read.search import global_index_live_row_filter -from pypaimon.utils.range import Range -from pypaimon.utils.roaring_bitmap import RoaringBitmap64 +from pypaimon.ray.search_execution import RaySearchTaskExecutor def _execute_full_text_search(builder, *, concurrency=None, ray_remote_args=None): - concurrency, remote_args = _execution_options(concurrency, ray_remote_args) - from pypaimon.table.source.search_dispatch import python_search_reader - reader = python_search_reader(builder.new_full_text_read()) - if (type(reader) is not DataEvolutionFullTextRead - or not reader._table.options.data_evolution_enabled()): - raise ValueError("Ray full-text search supports only data-evolution tables.") - _require_ray() + executor = RaySearchTaskExecutor(concurrency, ray_remote_args) + reader = builder.new_full_text_read() plan = builder.new_full_text_scan().scan() - return _RayFullTextRead(reader, concurrency, remote_args).read_plan(plan) - - -class _RayFullTextRead(DataEvolutionFullTextRead): - def __init__(self, reader, concurrency, remote_args): - super().__init__(reader._table, reader._limit, reader._text_columns, reader._query, - reader._partition_filter, reader._filter) - self._concurrency = concurrency - self._remote_args = remote_args - - def _worker_reader(self): - reader = DataEvolutionFullTextRead( - self._table.copy({"global-index.enabled": "false"}), self._limit, self._text_columns, self._query, - self._partition_filter, self._filter) - reader._snapshot = self._snapshot - return reader - - def read(self, splits): - # Keep bitmap construction and candidate data reads off the driver. - index_splits, raw_splits = _split_search_splits(splits) - indexed = self._eval_column_query({self._text_columns[0].name: index_splits}, None) - raw = self._read_raw_search(_raw_row_ranges(raw_splits), _index_type(index_splits)) - return indexed.or_(raw).top_k(self._limit) - - def _eval_column_query(self, splits_by_column, live_rows): - items = splits_by_column.get(self._text_columns[0].name, []) - merged = {} - reader = self._worker_reader() - reader._limit = _index_candidate_limit(items, self._limit) - with _shared_scalar_filters(items, self._filter, self._remote_args, - _prepare_scalar_filter) as prepare, closing(_map_tasks( - _search_index, reader, items, - self._concurrency, self._remote_args, True, prepare)) as tasks: - for _, scores in tasks: - # Duplicate row IDs keep the first planned shard's score. - for row_id, score in scores.items(): - merged.setdefault(row_id, score) - return DictBasedScoredIndexResult(merged).top_k(self._limit) - - def _read_raw_search(self, ranges, index_type, include_row_ids=None): - if not ranges: - return DictBasedScoredIndexResult({}) - include = None if include_row_ids is None else include_row_ids.serialize() - # Splitting this corpus would change document frequencies and BM25. - # Keep the existing local corpus intact, but read/build/search on a worker. - with closing(_map_tasks(_search_raw, self._worker_reader(), [(ranges, index_type, include)], - 1, self._remote_args)) as tasks: - for _, scores in tasks: - return DictBasedScoredIndexResult(scores) - - -def _search_index(reader, split, scalar_filter=None): - live = global_index_live_row_filter.live_rows( - reader._table, reader._partition_filter, reader._snapshot, - row_ranges=[Range(split.row_range_start, split.row_range_end)]) - if scalar_filter is None: - include = reader._matched_indexed_rows([split], live) - else: - bitmap, exact = scalar_filter - result = None if bitmap is None else GlobalIndexResult.create( - RoaringBitmap64.deserialize(bitmap), is_exact=exact) - include = reader._matched_indexed_rows([split], live, result) - if include is not None and include.is_empty(): - return {} - return _scores(reader._eval(split.row_range_start, split.row_range_end, - split.full_text_index_files, include).result()) - - -def _prepare_scalar_filter(reader, split): - result = reader._scalar_index_result(split.scalar_index_files, reader._snapshot) - return (None, False) if result is None else (result.results().serialize(), result.is_exact()) - - -def _search_raw(reader, item): - ranges, index_type, include = item - include = None if include is None else RoaringBitmap64.deserialize(include) - if reader._filter is not None: - live = global_index_live_row_filter.live_rows( - reader._table, reader._partition_filter, reader._snapshot, row_ranges=ranges) - candidates = GlobalIndexResult.from_ranges(ranges).results() - if live is not None: - candidates = RoaringBitmap64.and_(candidates, live) - if include is not None: - candidates = RoaringBitmap64.and_(candidates, include) - include = reader._matching_rows(candidates) - return _scores(reader._read_raw_search(ranges, index_type, include)) + return reader.read_plan(plan, executor=executor) def _execute_hybrid_search(builder, *, concurrency=None, ray_remote_args=None): - from pypaimon.ray.vector_search import _execute_vector_search, _task_cancellation - - concurrency, remote_args = _execution_options(concurrency, ray_remote_args) - # Resolve the shared snapshot before any route starts planning on a thread. - routes = builder.route_builders() - workers = min(concurrency, len(routes)) - route_concurrency = concurrency // max(1, workers) - - def execute_route(route): - execute = _execute_vector_search if route.route.is_vector() else _execute_full_text_search - return execute(route.search_builder, concurrency=route_concurrency, ray_remote_args=remote_args) - - if workers <= 1: - return builder.rank([builder.to_route_result(route, execute_route(route)) for route in routes]) - - cancelled = Event() - - def run(route): - with _task_cancellation(cancelled): - return execute_route(route) - - route_results = [None] * len(routes) - # Each concurrent route gets a share of the existing search-task budget. - # Completion order must not change route ordering supplied to the ranker. - with ThreadPoolExecutor(max_workers=workers) as executor: - futures = {} - try: - for ordinal, route in enumerate(routes): - futures[executor.submit(run, route)] = ordinal - for future in as_completed(futures): - ordinal = futures[future] - route_results[ordinal] = builder.to_route_result(routes[ordinal], future.result()) - finally: - cancelled.set() - for future in futures: - future.cancel() - return builder.rank(route_results) + return builder.execute(executor=RaySearchTaskExecutor(concurrency, ray_remote_args)) diff --git a/paimon-python/pypaimon/ray/search_execution.py b/paimon-python/pypaimon/ray/search_execution.py new file mode 100644 index 000000000000..91a4f1f64399 --- /dev/null +++ b/paimon-python/pypaimon/ray/search_execution.py @@ -0,0 +1,201 @@ +# 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. + +"""Bounded Ray execution of serializable backend search tasks.""" + +from concurrent.futures import CancelledError +from contextlib import contextmanager +import threading + +from pypaimon.table.source.search_execution import SearchTask, SearchTaskExecutor + + +def execution_options(concurrency, ray_remote_args): + concurrency = 4 if concurrency is None else concurrency + if isinstance(concurrency, bool) or not isinstance(concurrency, int) or concurrency < 1: + raise ValueError("concurrency must be a positive integer.") + remote_args = dict(ray_remote_args or {}) + if "num_returns" in remote_args: + raise ValueError("Vector search manages num_returns; omit it from ray_remote_args.") + return concurrency, remote_args + + +def require_ray(): + try: + import ray # noqa: F401 + except ModuleNotFoundError as error: + if error.name != "ray": + raise + raise ImportError("Ray vector search requires pypaimon[ray].") from error + + +def _execute_task(operation, context, item, *dependencies): + return SearchTask(operation, context, item).execute(*dependencies) + + +class _TaskBudget: + """Share in-flight slots across concurrent maps and hybrid route children.""" + + def __init__(self, concurrency): + self._slots = threading.BoundedSemaphore(concurrency) + self._condition = threading.Condition() + + def acquire(self): + return self._slots.acquire(blocking=False) + + def release(self): + self._slots.release() + with self._condition: + self._condition.notify_all() + + def wait(self): + with self._condition: + self._condition.wait(timeout=0.05) + + +class _CancellationState: + """Share query events across callback threads without sharing other queries.""" + + def __init__(self, events=()): + self._lock = threading.Lock() + self._events = dict.fromkeys(events, 1) + + @contextmanager + def cancellation_scope(self, event): + if event.is_set(): + raise CancelledError() + with self._lock: + self._events[event] = self._events.get(event, 0) + 1 + try: + yield + finally: + with self._lock: + count = self._events[event] - 1 + if count: + self._events[event] = count + else: + del self._events[event] + + def cancellation_events(self): + with self._lock: + return tuple(self._events) + + +class RaySearchTaskExecutor(SearchTaskExecutor): + """Run opaque backend tasks without constructing search readers.""" + + def __init__(self, concurrency=None, ray_remote_args=None, *, _budget=None, _events=()): + self._concurrency, self._remote_args = execution_options(concurrency, ray_remote_args) + self._budget = _TaskBudget(self._concurrency) if _budget is None else _budget + self._cancellations = _CancellationState(_events) + require_ray() + + @property + def concurrency(self): + return self._concurrency + + def with_concurrency(self, concurrency): + return RaySearchTaskExecutor(concurrency, self._remote_args, _budget=self._budget, + _events=self._cancellations.cancellation_events()) + + @contextmanager + def cancellation_scope(self, event): + # Native callbacks can enter map on a different driver thread. The + # query view owns active scopes; its children inherit event identities. + with self._cancellations.cancellation_scope(event): + yield + + def map(self, tasks, ordered=False): + """Bound pending plus buffered results, preserving optional input order.""" + import ray + + remaining = iter(enumerate(tasks)) + remote = ray.remote(_execute_task).options(**self._remote_args) + contexts = {} + dependencies = {} + pending = {} + buffered = {} + next_ordinal = 0 + exhausted = False + owned_slots = 0 + cancelled = self._cancellations.cancellation_events() + + def submit(task): + identity = id(task.context) + if identity not in contexts: + # Keep the object as well: context identity must not be reused. + contexts[identity] = (task.context, ray.put(task.context)) + refs = [] + for dependency in task.dependencies: + if dependency not in dependencies: + dependencies[dependency] = submit(dependency) + refs.append(dependencies[dependency]) + # Top-level dependencies resolve before Ray allocates the worker. + return remote.remote(task.operation, contexts[identity][1], task.item, *refs) + + try: + while True: + if any(event.is_set() for event in cancelled): + raise CancelledError() + while not exhausted and len(pending) + len(buffered) < self._concurrency: + if not self._budget.acquire(): + break + owned_slots += 1 + item = next(remaining, None) + if item is None: + exhausted = True + self._budget.release() + owned_slots -= 1 + break + ordinal, task = item + pending[submit(task)] = ordinal + if not pending: + if exhausted: + break + self._budget.wait() + continue + ready, _ = ray.wait(list(pending), num_returns=1, + **({"timeout": 0.1} if cancelled else {})) + if not ready: + continue + ref = ready[0] + result = ray.get(ref) + ordinal = pending.pop(ref) + if ordered: + buffered[ordinal] = result + del result + while next_ordinal in buffered: + result = buffered.pop(next_ordinal) + self._budget.release() + owned_slots -= 1 + yield next_ordinal, result + del result + next_ordinal += 1 + else: + self._budget.release() + owned_slots -= 1 + yield ordinal, result + del result + finally: + for ref in list(pending) + list(dependencies.values()): + try: + ray.cancel(ref) + except Exception: + # Preserve the primary failure when a cluster has stopped. + pass + for _ in range(owned_slots): + self._budget.release() diff --git a/paimon-python/pypaimon/ray/vector_index_build.py b/paimon-python/pypaimon/ray/vector_index_build.py index b0c4d1c7a2e6..31bf1a73a5a3 100644 --- a/paimon-python/pypaimon/ray/vector_index_build.py +++ b/paimon-python/pypaimon/ray/vector_index_build.py @@ -20,13 +20,13 @@ from pypaimon.execution.python.globalindex.build_plan import split_by_global_index_shard from pypaimon.execution.python.globalindex.index_file_utils import new_global_index_file_name from pypaimon.globalindex.vindex_options import VINDEX_IDENTIFIERS -from pypaimon.ray.vector_search import _execution_options +from pypaimon.ray.search_execution import execution_options def validate_build_options(builder, concurrency, ray_remote_args): if builder._index_type not in VINDEX_IDENTIFIERS: raise ValueError("Ray index building supports only native vector indexes.") - concurrency, options = _execution_options(concurrency, ray_remote_args) + concurrency, options = execution_options(concurrency, ray_remote_args) # Retrying a write task can race with a lost worker writing the same file. # A failed build can instead be rerun after its uncommitted outputs close. if options.get("max_retries", 0) != 0 or options.get("retry_exceptions", False): diff --git a/paimon-python/pypaimon/ray/vector_search.py b/paimon-python/pypaimon/ray/vector_search.py index 5a326d7946e7..6400bb2ef105 100644 --- a/paimon-python/pypaimon/ray/vector_search.py +++ b/paimon-python/pypaimon/ray/vector_search.py @@ -15,272 +15,13 @@ # specific language governing permissions and limitations # under the License. -"""Ray execution of snapshot-pinned, single-vector queries.""" +"""Execute vector plans with the public search-task protocol and Ray scheduler.""" -from collections import Counter -from concurrent.futures import CancelledError -from contextlib import closing, contextmanager -import math -import threading - -from pypaimon.globalindex.global_index_result import GlobalIndexResult -from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult -from pypaimon.execution.python.read.search.vector_search_read import ( - DataEvolutionVectorRead, - _filtered_raw_row_ranges, - _offer_score, - _scored_result, - _split_search_splits, -) -from pypaimon.execution.python.read.search import global_index_live_row_filter -from pypaimon.utils.range import Range -from pypaimon.utils.roaring_bitmap import RoaringBitmap64 +from pypaimon.ray.search_execution import RaySearchTaskExecutor def _execute_vector_search(builder, *, concurrency=None, ray_remote_args=None): - """Execute the builder of an already snapshot-pinned ``VectorQuery``.""" - concurrency, remote_args = _execution_options(concurrency, ray_remote_args) - from pypaimon.table.source.search_dispatch import python_search_reader - reader = python_search_reader(builder.new_vector_search_read()) - if (type(reader) is not DataEvolutionVectorRead - or not reader._table.options.data_evolution_enabled()): - raise ValueError("Ray vector search supports only data-evolution tables.") - if not all(math.isfinite(float(value)) for value in reader._query_vector): - raise ValueError("Ray vector search requires a finite query vector.") - - _require_ray() + executor = RaySearchTaskExecutor(concurrency, ray_remote_args) + reader = builder.new_vector_search_read() plan = builder.new_vector_search_scan().scan() - return _RayVectorSearchRead(reader, concurrency, remote_args).read_plan(plan) - - -def _execution_options(concurrency, ray_remote_args): - concurrency = 4 if concurrency is None else concurrency - if isinstance(concurrency, bool) or not isinstance(concurrency, int) or concurrency < 1: - raise ValueError("concurrency must be a positive integer.") - remote_args = dict(ray_remote_args or {}) - if "num_returns" in remote_args: - raise ValueError("Vector search manages num_returns; omit it from ray_remote_args.") - return concurrency, remote_args - - -def _require_ray(): - try: - import ray # noqa: F401 - except ModuleNotFoundError as error: - if error.name != "ray": - raise - raise ImportError("Ray vector search requires pypaimon[ray].") from error - - -class _RayVectorSearchRead(DataEvolutionVectorRead): - """Replace split execution while retaining local merge/refinement semantics.""" - - def __init__(self, reader, concurrency, remote_args): - super().__init__( - reader._table, reader._limit, reader._vector_column, reader._query_vector, - reader._filter, reader._partition_filter, reader._options) - self._concurrency = concurrency - self._remote_args = remote_args - - def _prepare_search_splits(self, splits, snapshot): - # Scalar evaluation and candidate verification belong to the workers. - return _split_search_splits(splits) - - def _search_index_splits(self, splits, query, search_limit, snapshot, batch=False): - # Send predicates and metadata, leaving each worker to build and consume - # its own filter. Batch queries reuse the same filter for every vector. - context = (self._table, self._vector_column, query, search_limit, self._options, - self._filter, self._partition_filter, snapshot) - with _shared_scalar_filters(splits, self._filter, self._remote_args) as prepare, closing(_map_tasks( - _search_index_split, context, splits, self._concurrency, self._remote_args, True, prepare)) as tasks: - for _, (metric, scores) in tasks: - if metric is not None: - self._set_index_metric(metric) - # Keep plan order, including duplicate-row precedence, regardless - # of worker completion order. Do not refine or truncate per worker. - yield DictBasedScoredIndexResult(scores) - del scores - - def _read_raw_search(self, raw_row_ranges, pre_filter, query_vector, - index_type=None, include_filter=True, - score_candidates=None, snapshot=None): - ranges = _filtered_raw_row_ranges(raw_row_ranges, pre_filter) - if score_candidates is not None: - # Refine only the globally selected candidates. Intersect before - # planning so workers neither read nor score other rows. - ranges = _filtered_raw_row_ranges(ranges, score_candidates.to_range_list()) - if not ranges: - return DictBasedScoredIndexResult({}) - table_read, splits = self._plan_raw_read(ranges, include_filter, snapshot) - context = (table_read, self._vector_column, query_vector, - self._limit, self._search_metric(index_type)) - heap = [] - # Keep each planned data split intact: its files may be vertical column - # groups that must be read together to reconstruct a logical row. - with closing(_map_tasks( - _search_raw_split, context, splits, self._concurrency, self._remote_args)) as tasks: - for _, scores in tasks: - for row_id, score in scores.items(): - _offer_score(heap, self._limit, row_id, score) - return _scored_result(heap) - - -def _search_index_split(context, split, scalar_filter=None): - table, column, query, limit, options, filter_, partition_filter, snapshot = context - # Explicit shard indexes are evaluated below. Candidate verification must - # not rediscover and scan table-wide scalar indexes in every worker. - table = table.copy({"global-index.enabled": "false"}) - reader = DataEvolutionVectorRead( - table, limit, column, query, filter_, partition_filter, options) - include = _worker_pre_filter(reader, split, snapshot, scalar_filter) - result = reader._eval( - split.row_range_start, split.row_range_end, - split.vector_index_files, query, limit, include).result() - # The reader closes its native resources before returning a completed result. - # Return metric even for zero hits, so inconsistent indexes cannot be hidden - # by a scalar filter and raw fallback uses the persisted index metric. - return reader._index_metric, _scores(result) - - -def _worker_pre_filter(reader, split, snapshot, scalar_filter): - live_rows = global_index_live_row_filter.live_rows( - reader._table, reader._partition_filter, snapshot, - row_ranges=[Range(split.row_range_start, split.row_range_end)]) - if scalar_filter is None: - matched = reader._scalar_matched_rows([split], snapshot) - else: - bitmap, exact = scalar_filter - result = None if bitmap is None else GlobalIndexResult.create( - RoaringBitmap64.deserialize(bitmap), is_exact=exact) - matched = reader._refine_scalar_result([split], result, snapshot) - filters = reader._combine_pre_filters([split], live_rows, matched) - return filters[0] if filters else None - - -def _prepare_scalar_filter(context, split): - table, column, query, limit, options, filter_, partition_filter, snapshot = context - reader = DataEvolutionVectorRead(table, limit, column, query, filter_, partition_filter, options) - result = reader._scalar_index_result([split], snapshot) - return (None, False) if result is None else (result.results().serialize(), result.is_exact()) - - -@contextmanager -def _shared_scalar_filters(splits, predicate, remote_args, worker=None): - """Evaluate identical scalar-index inputs once, without fetching masks on the driver.""" - import ray - - def key(split): - return tuple(sorted(f.file_name for f in split.scalar_index_files)) - - counts = Counter(key(split) for split in splits) if predicate is not None else {} - shared = {files for files, count in counts.items() if files and count > 1} - pending = {} - remote = ray.remote(worker or _prepare_scalar_filter).options(**remote_args) if shared else None - - def prepare(context_ref, split): - files = key(split) if shared else None - if files not in shared: - return split, None - if files not in pending: - pending[files] = remote.remote(context_ref, split) - # A top-level ObjectRef dependency is resolved before allocating the - # search worker, so waiting searches cannot starve their filter tasks. - return split, pending[files] - - try: - yield prepare - finally: - for ref in pending.values(): - try: - ray.cancel(ref) - except Exception: - # Preserve the search failure if the cluster is already gone. - pass - - -def _search_raw_split(context, split): - from pypaimon.read.read_support import _ClosableArrowBatchReader - - table_read, column, query, limit, metric = context - scorer = DataEvolutionVectorRead(table_read.table, limit, column, query) - reader, batches = table_read._new_arrow_batch_reader([split]) - with _ClosableArrowBatchReader(reader, batches) as batch_reader: - return _scores(scorer._score_raw_batches(batch_reader, query, metric, reject_nan=True)) - - -def _scores(result): - if result is None: - return {} - getter = result.score_getter() - scores = {row_id: getter(row_id) for row_id in result.results()} - if any(math.isnan(score) for score in scores.values()): - raise ValueError("Ray vector search cannot rank NaN scores.") - return scores - - -_task_state = threading.local() - - -@contextmanager -def _task_cancellation(event): - """Keep hybrid cancellation on driver threads, out of serialized worker state.""" - previous = getattr(_task_state, "cancelled", None) - _task_state.cancelled = event - try: - if event.is_set(): - raise CancelledError() - yield - finally: - _task_state.cancelled = previous - - -def _map_tasks(worker, context, items, concurrency, remote_args, ordered=False, prepare_item=None): - """Bound in-flight work and, optionally, completed results awaiting plan order.""" - import ray - - if not items: - return - # Share only serializable planning state, never open readers or closures. - context_ref = ray.put(context) - remote = ray.remote(worker).options(**remote_args) - remaining = iter(enumerate(items)) - pending = {} - buffered = {} - next_ordinal = 0 - cancelled = getattr(_task_state, "cancelled", None) - try: - while True: - if cancelled is not None and cancelled.is_set(): - raise CancelledError() - while len(pending) + len(buffered) < concurrency: - item = next(remaining, None) - if item is None: - break - ordinal, split = item - args = (split,) if prepare_item is None else prepare_item(context_ref, split) - pending[remote.remote(context_ref, *args)] = ordinal - if not pending: - break - ready, _ = ray.wait(list(pending), num_returns=1, - **({"timeout": 0.1} if cancelled is not None else {})) - if not ready: - continue - ref = ready[0] - result = ray.get(ref) - ordinal = pending.pop(ref) - if ordered: - buffered[ordinal] = result - del result - while next_ordinal in buffered: - yield next_ordinal, buffered.pop(next_ordinal) - next_ordinal += 1 - else: - yield ordinal, result - del result - finally: - for ref in pending: - try: - ray.cancel(ref) - except Exception: - # Preserve the search failure if the cluster is already gone. - pass + return reader.read_plan(plan, executor=executor) diff --git a/paimon-python/pypaimon/table/source/hybrid_search_builder.py b/paimon-python/pypaimon/table/source/hybrid_search_builder.py index 27a1d2fff4fe..043765dc0ce2 100644 --- a/paimon-python/pypaimon/table/source/hybrid_search_builder.py +++ b/paimon-python/pypaimon/table/source/hybrid_search_builder.py @@ -228,6 +228,13 @@ def execute_local(self) -> ScoredGlobalIndexResult: from pypaimon.execution.python.read.search.hybrid_search import execute_local return execute_local(self) + def execute(self, executor=None) -> ScoredGlobalIndexResult: + """Execute complete routes with an optional search-task scheduler.""" + if executor is None: + return self.execute_local() + from pypaimon.execution.python.read.search.hybrid_search import execute + return execute(self, executor) + class HybridSearchBuilderImpl(HybridSearchBuilder): """Implementation for HybridSearchBuilder.""" @@ -244,6 +251,12 @@ def execute_local(self) -> ScoredGlobalIndexResult: from pypaimon.table.source.search_dispatch import execute_hybrid_search return execute_hybrid_search(self) + def execute(self, executor=None) -> ScoredGlobalIndexResult: + if executor is None: + return self.execute_local() + from pypaimon.table.source.search_dispatch import execute_hybrid_search + return execute_hybrid_search(self, executor) + def with_limit(self, limit: int) -> 'HybridSearchBuilder': self._limit = limit return self diff --git a/paimon-python/pypaimon/table/source/search_dispatch.py b/paimon-python/pypaimon/table/source/search_dispatch.py index 6589a3fe8e9e..14446063941d 100644 --- a/paimon-python/pypaimon/table/source/search_dispatch.py +++ b/paimon-python/pypaimon/table/source/search_dispatch.py @@ -55,17 +55,21 @@ def _reader_for_plan(self, plan): reader._table = reader._table._copy_with_snapshot(plan.snapshot()) return reader - def read_plan(self, plan): + def read_plan(self, plan, executor=None): native = self._native_reader.resolve_plan(plan) if native is not None: - return self._native_reader.read_plan(native) - return self._reader_for_plan(plan).read_plan(plan) + return (self._native_reader.read_plan(native) if executor is None else + self._native_reader.read_plan(native, executor=executor)) + reader = self._reader_for_plan(plan) + return reader.read_plan(plan) if executor is None else reader.read_plan(plan, executor=executor) - def read_batch_plan(self, plan): + def read_batch_plan(self, plan, executor=None): native = self._native_reader.resolve_plan(plan) if native is not None: - return self._native_reader.read_batch_plan(native) - return self._reader_for_plan(plan).read_batch_plan(plan) + return (self._native_reader.read_batch_plan(native) if executor is None else + self._native_reader.read_batch_plan(native, executor=executor)) + reader = self._reader_for_plan(plan) + return reader.read_batch_plan(plan) if executor is None else reader.read_batch_plan(plan, executor=executor) class _SplitSearchRead(_SearchRead): @@ -79,7 +83,7 @@ def read_batch(self, splits): def python_search_reader(reader): - """Select Python explicitly for Ray split workers or plan comparisons.""" + """Select Python explicitly for backend parity and transported-plan checks.""" return reader._python_reader() if isinstance(reader, _SearchRead) else reader @@ -144,11 +148,12 @@ def execute_full_text_search(builder): return builder.new_full_text_read().read_plan(builder.new_full_text_scan().scan()) -def execute_hybrid_search(builder): +def execute_hybrid_search(builder, executor=None): if _native_requested(builder): from pypaimon.execution.native.read.native_hybrid_search import try_native_hybrid_search - result = try_native_hybrid_search(builder) + result = (try_native_hybrid_search(builder) if executor is None else + try_native_hybrid_search(builder, executor=executor)) if result is not None: return result - from pypaimon.execution.python.read.search.hybrid_search import execute_local - return execute_local(builder) + from pypaimon.execution.python.read.search.hybrid_search import execute + return execute(builder, executor) diff --git a/paimon-python/pypaimon/table/source/search_execution.py b/paimon-python/pypaimon/table/source/search_execution.py new file mode 100644 index 000000000000..6da31b3f9560 --- /dev/null +++ b/paimon-python/pypaimon/table/source/search_execution.py @@ -0,0 +1,73 @@ +# 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. + +"""Transport and scheduling contracts for backend-owned search tasks.""" + +from abc import ABC, abstractmethod +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, Callable, Tuple + + +@dataclass(frozen=True, eq=False) +class SearchTask: + """A serializable operation with shared context and optional prerequisites. + + Dependencies use object identity within a map call. Executors resolve them + before scheduling their dependents, without occupying a dependent worker. + Context must contain transportable configuration, not an open reader. + """ + + operation: Callable + context: Any + item: Any + dependencies: Tuple['SearchTask', ...] = () + + def execute(self, *dependency_results): + return self.operation(self.context, self.item, *dependency_results) + + +class SearchTaskExecutor(ABC): + """Schedule search tasks while the selected backend coordinates the query.""" + + @abstractmethod + def map(self, tasks, ordered=False): + """Return a closable iterator of ``(input_ordinal, result)`` pairs. + + ``ordered`` requires input order, including duplicate-row precedence. + Closing the iterator or failing an operation cancels its pending work. + Exceptions propagate without replaying through another backend. + """ + pass + + @property + def concurrency(self): + """Number of concurrent tasks available to a query.""" + return 1 + + def with_concurrency(self, concurrency): + """Create a query view sharing resources and the active parent cancellation. + + Implementations with cancellation state isolate future scopes on the + returned view, so reusing an executor does not cancel unrelated queries. + """ + return self + + @contextmanager + def cancellation_scope(self, event): + """Observe a query's cooperative cancellation event when supported.""" + yield diff --git a/paimon-python/pypaimon/table/source/search_types.py b/paimon-python/pypaimon/table/source/search_types.py index a649be95fcc9..a744afa518c5 100644 --- a/paimon-python/pypaimon/table/source/search_types.py +++ b/paimon-python/pypaimon/table/source/search_types.py @@ -19,11 +19,12 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import List +from typing import List, Optional from pypaimon.globalindex.global_index_result import GlobalIndexResult from pypaimon.read.split import DataSplit from pypaimon.table.source.full_text_search_split import FullTextSearchSplit +from pypaimon.table.source.search_execution import SearchTaskExecutor class VectorSearchScanPlan: @@ -54,9 +55,14 @@ def scan(self): class VectorSearchRead(ABC): """Vector search read to read index files.""" - def read_plan(self, plan): - # type: (VectorSearchScanPlan) -> GlobalIndexResult - return self._read(plan.splits(), plan.snapshot()) + def read_plan(self, plan, executor=None): + # type: (VectorSearchScanPlan, Optional[SearchTaskExecutor]) -> GlobalIndexResult + """Execute the complete captured plan with an optional task scheduler.""" + reader = self if executor is None else self._with_executor(executor) + return reader._read(plan.splits(), plan.snapshot()) + + def _with_executor(self, executor): + raise NotImplementedError('This vector reader does not support task execution') def read(self, splits): # type: (List[VectorSearchSplit]) -> GlobalIndexResult @@ -70,9 +76,14 @@ def _read(self, splits, snapshot): class BatchVectorSearchRead(ABC): """Batch vector search read to read index files.""" - def read_batch_plan(self, plan): - # type: (VectorSearchScanPlan) -> List[GlobalIndexResult] - return self._read_batch(plan.splits(), plan.snapshot()) + def read_batch_plan(self, plan, executor=None): + # type: (VectorSearchScanPlan, Optional[SearchTaskExecutor]) -> List[GlobalIndexResult] + """Execute all queries over the captured plan, retaining input order.""" + reader = self if executor is None else self._with_executor(executor) + return reader._read_batch(plan.splits(), plan.snapshot()) + + def _with_executor(self, executor): + raise NotImplementedError('This batch vector reader does not support task execution') def read_batch(self, splits): # type: (List[VectorSearchSplit]) -> List[GlobalIndexResult] @@ -108,8 +119,13 @@ def scan(self) -> FullTextScanPlan: class FullTextRead(ABC): """Full-text read to read index files.""" - def read_plan(self, plan: FullTextScanPlan) -> GlobalIndexResult: - return self.read(plan.splits()) + def read_plan(self, plan: FullTextScanPlan, executor: Optional[SearchTaskExecutor] = None) -> GlobalIndexResult: + """Execute a complete plan, keeping corpus-wide ranking in the backend.""" + reader = self if executor is None else self._with_executor(executor) + return reader.read(plan.splits()) + + def _with_executor(self, executor): + raise NotImplementedError('This full-text reader does not support task execution') @abstractmethod def read(self, splits: List[FullTextSearchSplit]) -> GlobalIndexResult: diff --git a/paimon-python/pypaimon/tests/native_search_plan_test.py b/paimon-python/pypaimon/tests/native_search_plan_test.py index 19b9bb89b2e8..0d43e591d3c3 100644 --- a/paimon-python/pypaimon/tests/native_search_plan_test.py +++ b/paimon-python/pypaimon/tests/native_search_plan_test.py @@ -17,7 +17,7 @@ """Explicit search plans keep native work and transportable source metadata.""" -from contextlib import ExitStack, contextmanager +from contextlib import ExitStack from copy import copy import pickle from unittest.mock import patch @@ -327,48 +327,34 @@ def test_native_plan_is_query_independent_and_pickle_keeps_original_handle(table @pytest.mark.parametrize('family', ['vector', 'batch', 'full-text']) def test_ray_accepts_native_reader_and_transports_captured_sources(table, family): - # Run the real Ray adapters and workers with in-process serialized task - # dispatch. This exercises their reader qualification and metadata contract - # without starting another cluster in the native CI jobs. + # Exercise the real Native task bridge with process-style serialization, + # without rebuilding the SDK or starting another cluster in Native CI. from pypaimon.ray import vector_search, batch_vector_search, full_text_search - module, reader_type, execute, method = { - 'vector': (vector_search, vector_search._RayVectorSearchRead, - vector_search._execute_vector_search, 'read_plan'), - 'batch': (batch_vector_search, batch_vector_search._RayBatchVectorSearchRead, - batch_vector_search._execute_batch_vector_search, 'read_batch_plan'), - 'full-text': (full_text_search, full_text_search._RayFullTextRead, - full_text_search._execute_full_text_search, 'read_plan'), + from pypaimon.table.source.search_execution import SearchTaskExecutor + module, execute = { + 'vector': (vector_search, vector_search._execute_vector_search), + 'batch': (batch_vector_search, batch_vector_search._execute_batch_vector_search), + 'full-text': (full_text_search, full_text_search._execute_full_text_search), }[family] classic = _builder(table, family, limit=20) expected = _read(_reader(classic, family), family, _scan(classic, family).scan()) - original_read = getattr(reader_type, method) - - def read_after_append(reader, plan): - captured = plan.snapshot().id - transported = pickle.loads(pickle.dumps(plan)) - assert plan._unchanged_native_plan() is not None - indexes._append(table, hybrid._rows(12, 4), hybrid.SCHEMA) - assert transported.snapshot().id == captured - return original_read(reader, transported) - - def dispatch(worker, context, items, *args): - for ordinal, item in enumerate(items): - yield ordinal, worker(pickle.loads(pickle.dumps(context)), pickle.loads(pickle.dumps(item))) - - @contextmanager - def scalar_filters(*args): - yield None - - with ExitStack() as stack: - stack.enter_context(patch.object(module, '_require_ray')) - stack.enter_context(patch.object(module, '_map_tasks', dispatch)) - stack.enter_context(patch.object(reader_type, method, read_after_append)) - if family != 'full-text': - stack.enter_context(patch.object(module, '_shared_scalar_filters', scalar_filters)) - for cls in (DataEvolutionVectorScan, DataEvolutionFullTextScan): - stack.enter_context(patch.object(cls, 'scan', side_effect=AssertionError('Python search scan ran'))) + appended = [] + + class SerializedExecutor(SearchTaskExecutor): + def map(self, tasks, ordered=False): + assert ordered + if not appended: + indexes._append(table, hybrid._rows(12, 4), hybrid.SCHEMA) + appended.append(True) + for ordinal, task in enumerate(tasks): + transported = pickle.loads(pickle.dumps(task)) + assert not transported.dependencies + yield ordinal, transported.execute() + + with _native_only(), patch.object(module, 'RaySearchTaskExecutor', return_value=SerializedExecutor()): actual = execute(_builder(table, family, native=True, limit=20)) _assert_results(actual if family == 'batch' else [actual], expected) + assert appended @pytest.mark.parametrize('family', ['vector', 'full-text']) diff --git a/paimon-python/pypaimon/tests/ray_batch_vector_refinement_test.py b/paimon-python/pypaimon/tests/ray_batch_vector_refinement_test.py index f449eb3aadb9..3edb25301fd6 100644 --- a/paimon-python/pypaimon/tests/ray_batch_vector_refinement_test.py +++ b/paimon-python/pypaimon/tests/ray_batch_vector_refinement_test.py @@ -23,8 +23,11 @@ pytest.importorskip("ray") +from pypaimon.ray.search_execution import RaySearchTaskExecutor +from pypaimon.tests.search_execution_test_helpers import task_parts + from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult -from pypaimon.ray import batch_vector_search as search_module +from pypaimon.execution.python.read.search import distributed_vector_search as search_module from pypaimon.read.table_read import TableRead from pypaimon.tests import ray_vector_search_test as fixtures from pypaimon.tests.ray_vector_search_test import VECTORS, add_rows, build_index @@ -44,20 +47,21 @@ def test_batch_refinement_reads_vectors_only_on_workers(table, ray_cluster, metr query = table.search_vectors(QUERIES, options={"refine_factor": "2"}).select(["id"]).limit(3) expected = query.to_list() original_read = TableRead._new_arrow_batch_reader - original_map = search_module._map_tasks + original_map = RaySearchTaskExecutor.map dispatched = [] def no_driver_vectors(read, *args, **kwargs): assert "embedding" not in [field.name for field in read.read_type] return original_read(read, *args, **kwargs) - def record(worker, context, items, *args): + def record(executor, tasks, ordered=False): + worker, context, items = task_parts(tasks) if worker is search_module._search_batch_refine_split: dispatched.append(len(items)) - return original_map(worker, context, items, *args) + return original_map(executor, tasks, ordered=ordered) with patch.object(TableRead, "_new_arrow_batch_reader", no_driver_vectors), \ - patch.object(search_module, "_map_tasks", record): + patch.object(RaySearchTaskExecutor, "map", record): actual = query.to_arrow(execution="ray", concurrency=concurrency) assert [result.to_pylist() for result in actual] == expected assert dispatched == [2] @@ -99,9 +103,9 @@ def batches(): def test_empty_candidates_do_not_plan_or_dispatch(table): query = table.search_vectors(QUERIES, options={"refine_factor": "2"}) reader = query._batch_vector_search_builder(query).new_batch_vector_search_read() - distributed = search_module._RayBatchVectorSearchRead(reader, 2, {}) + distributed = search_module.DistributedBatchVectorSearchRead(reader, RaySearchTaskExecutor(2, {})) with patch.object(distributed, "_plan_raw_read", side_effect=AssertionError("empty read")), \ - patch.object(search_module, "_map_tasks", side_effect=AssertionError("empty task")): + patch.object(RaySearchTaskExecutor, "map", side_effect=AssertionError("empty task")): results = distributed._maybe_rerank_indexed_results( [DictBasedScoredIndexResult({}) for _ in QUERIES], "ivf-flat", QUERIES) assert all(result.results().is_empty() for result in results) @@ -113,16 +117,17 @@ def test_commit_between_selection_and_refinement_keeps_snapshot(table, ray_clust table.raw_table = table.raw_table.copy({"global-index.column-update-action": "DROP_PARTITION_INDEX"}) query = table.search_vectors(QUERIES, options={"refine_factor": "2"}).select(["id"]).limit(1) expected = query.to_list() - original = search_module._map_tasks + original = RaySearchTaskExecutor.map changed = [] - def commit(worker, *args): + def commit(executor, tasks, ordered=False): + worker, context, items = task_parts(tasks) if worker is search_module._search_batch_refine_split and not changed: table.update("id = 0", {"embedding": [100., 1.]}) changed.append(True) - return original(worker, *args) + return original(executor, tasks, ordered=ordered) - with patch.object(search_module, "_map_tasks", commit): + with patch.object(RaySearchTaskExecutor, "map", commit): assert [result.to_pylist() for result in query.to_arrow(execution="ray")] == expected assert changed == [True] assert query.to_list() != expected diff --git a/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py b/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py index 93045c56af1e..c2f04c716298 100644 --- a/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py +++ b/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py @@ -24,11 +24,14 @@ pytest.importorskip("ray") +from pypaimon.ray.search_execution import RaySearchTaskExecutor +from pypaimon.tests.search_execution_test_helpers import task_parts + import pypaimon.multimodal as pm from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult from pypaimon.index.index_file_meta import IndexFileMeta from pypaimon.multimodal.query import ScanQuery -from pypaimon.ray import batch_vector_search as search_module +from pypaimon.execution.python.read.search import distributed_vector_search as search_module from pypaimon.execution.python.read.search.vector_search_read import BatchVectorSearchReadImpl from pypaimon.execution.python.read.search.vector_search_scan import DataEvolutionVectorScan from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit @@ -57,15 +60,16 @@ def test_raw_batch_shares_split_tasks(table, ray_cluster, metric, concurrency): else: scores = vectors.dot(vector) expected.append(sorted(sorted(range(len(VECTORS)), key=lambda i: (-scores[i], i))[:3])) - original = search_module._map_tasks + original = RaySearchTaskExecutor.map dispatched = [] - def record(worker, context, items, *args): + def record(executor, tasks, ordered=False): + worker, context, items = task_parts(tasks) if worker is search_module._search_batch_raw_split: dispatched.append(len(items)) - return original(worker, context, items, *args) + return original(executor, tasks, ordered=ordered) - with patch.object(search_module, "_map_tasks", record): + with patch.object(RaySearchTaskExecutor, "map", record): actual = query.to_arrow(execution="ray", concurrency=concurrency) assert [ids(result) for result in actual] == expected assert dispatched == [2] @@ -80,14 +84,15 @@ def test_indexed_batch_and_raw_fallback_use_persisted_metric(table, ray_cluster, options = {"refine_factor": "2"} if refine else {} query = table.search_vectors(QUERIES, options=options).select(["id"]).limit(3) expected = query.to_list() - original = search_module._map_tasks + original = RaySearchTaskExecutor.map dispatched = [] - def record(worker, context, items, *args): + def record(executor, tasks, ordered=False): + worker, context, items = task_parts(tasks) dispatched.append((worker.__name__, len(items))) - return original(worker, context, items, *args) + return original(executor, tasks, ordered=ordered) - with patch.object(search_module, "_map_tasks", record): + with patch.object(RaySearchTaskExecutor, "map", record): actual = query.to_arrow(execution="ray", concurrency=2) assert [result.to_pylist() for result in actual] == expected assert ("_search_batch_index_split", 2) in dispatched @@ -186,16 +191,16 @@ def test_whole_batch_pins_snapshot_and_can_be_reused(table, ray_cluster, stage): add_rows(table, VECTORS) query = table.search_vectors(QUERIES).select(["id", "category"]).limit(1) expected = query.to_list() - target = search_module if stage == "dispatch" else ScanQuery - name = "_map_tasks" if stage == "dispatch" else "_read_global_index_result" + target = RaySearchTaskExecutor if stage == "dispatch" else ScanQuery + name = "map" if stage == "dispatch" else "_read_global_index_result" original = getattr(target, name) changed = [] - def commit(*args): + def commit(*args, **kwargs): if not changed: table.update("id = 0", {"category": "changed"}) changed.append(True) - return original(*args) + return original(*args, **kwargs) with patch.object(target, name, commit): assert [result.to_pylist() for result in query.to_arrow(execution="ray")] == expected @@ -230,7 +235,7 @@ def append_before_scan(search_scan): return scan(search_scan) with patch.object(DataEvolutionVectorScan, "scan", append_before_scan), \ - patch.object(search_module, "_map_tasks", side_effect=AssertionError("empty task")): + patch.object(RaySearchTaskExecutor, "map", side_effect=AssertionError("empty task")): results = query.to_arrow(execution="ray") assert len(results) == 3 assert all(result.schema.names == ["id"] and result.num_rows == 0 for result in results) @@ -241,22 +246,23 @@ def test_global_candidates_and_duplicate_precedence(table): add_rows(table, [[0., 1.], [10., 1.], [20., 1.], [1., 1.]]) query = table.search_vectors([[1., 1.], [20., 1.]], options={"refine_factor": "2"}).limit(1)._for_execution() reader = query._batch_vector_search_builder(query).new_batch_vector_search_read() - distributed = search_module._RayBatchVectorSearchRead(reader, 2, {}) + distributed = search_module.DistributedBatchVectorSearchRead(reader, RaySearchTaskExecutor(2, {})) index_file = IndexFileMeta("ivf-flat", "fake", 1, 2) splits = [IndexVectorSearchSplit(0, 1, [index_file]), IndexVectorSearchSplit(0, 3, [index_file])] - def completed(worker, context, items, *args): + def completed(executor, tasks, ordered=False): + worker, context, items = task_parts(tasks) if worker is search_module._search_batch_refine_split: for ordinal, split in enumerate(items): yield ordinal, worker(context, split) return # First query must exclude exact nearest row 3; second query has its own candidates. # Duplicate row 0 must retain the earlier split's score for the second query. - assert args[2] is True + assert ordered is True yield 0, ("l2", [{0: 10., 1: 9.}, {0: 1., 1: 9.}]) yield 1, ("l2", [{2: 8., 3: 7.}, {0: 100., 2: 8., 3: 7.}]) - with patch.object(search_module, "_map_tasks", completed): + with patch.object(RaySearchTaskExecutor, "map", completed): results = distributed._read_batch(splits, query._table._read_snapshot) assert [list(result.results()) for result in results] == [[0], [2]] assert results[0].score_getter()(0) == 0.5 @@ -265,17 +271,17 @@ def completed(worker, context, items, *args): def test_metric_mismatch_in_empty_shard_closes_tasks(table): query = table.search_vectors(QUERIES) reader = query._batch_vector_search_builder(query).new_batch_vector_search_read() - distributed = search_module._RayBatchVectorSearchRead(reader, 2, {}) + distributed = search_module.DistributedBatchVectorSearchRead(reader, RaySearchTaskExecutor(2, {})) closed = [] - def completed(*args): + def completed(executor, tasks, ordered=False): try: yield 0, ("l2", [{}, {}, {}]) yield 1, ("inner_product", [{1: 10.}] * 3) finally: closed.append(True) - with patch.object(search_module, "_map_tasks", completed), \ + with patch.object(RaySearchTaskExecutor, "map", completed), \ pytest.raises(ValueError, match="different metrics"): list(distributed._search_index_splits([None, None], QUERIES, 1, None, batch=True)) assert closed == [True] diff --git a/paimon-python/pypaimon/tests/ray_full_text_search_test.py b/paimon-python/pypaimon/tests/ray_full_text_search_test.py index c2521a4a3a32..42199e809b81 100644 --- a/paimon-python/pypaimon/tests/ray_full_text_search_test.py +++ b/paimon-python/pypaimon/tests/ray_full_text_search_test.py @@ -20,12 +20,15 @@ import pytest pytest.importorskip("ray") + +from pypaimon.ray.search_execution import RaySearchTaskExecutor +from pypaimon.tests.search_execution_test_helpers import task_parts pytest.importorskip("paimon_ftindex") import pypaimon.multimodal as pm from pypaimon.tests import ray_vector_search_test as ray_fixtures from pypaimon.tests import full_text_scalar_filter_test as text_fixtures -from pypaimon.ray import full_text_search as module +from pypaimon.execution.python.read.search import distributed_full_text_search as module ray_cluster = ray_fixtures.ray_cluster docs = text_fixtures.docs @@ -39,13 +42,14 @@ def test_index_raw_filters_and_scores(docs, ray_cluster, predicate, concurrency) query = docs.search("paimon", column="text", pre_filter=predicate).select(["id"]).with_score().order_by_score() expected = query.to_list() dispatched = [] - original = module._map_tasks + original = RaySearchTaskExecutor.map - def record(worker, context, items, *args): + def record(executor, tasks, ordered=False): + worker, context, items = task_parts(tasks) dispatched.append((worker.__name__, len(items))) - return original(worker, context, items, *args) + return original(executor, tasks, ordered=ordered) - with patch.object(module, "_map_tasks", record): + with patch.object(RaySearchTaskExecutor, "map", record): assert query.to_list(execution="ray", concurrency=concurrency) == expected if predicate is None: assert ("_search_index", 1) in dispatched @@ -149,16 +153,16 @@ def record(reader, split): def test_worker_filters_use_planned_snapshot(docs, ray_cluster): query = docs.search("paimon", column="text", pre_filter="id >= 0").select(["id"]).order_by_score() expected = query.to_list() - original = module._map_tasks + original = RaySearchTaskExecutor.map committed = [] - def commit_before_dispatch(*args): + def commit_before_dispatch(executor, tasks, ordered=False): if not committed: docs.delete("id = 0") committed.append(True) - return original(*args) + return original(executor, tasks, ordered=ordered) - with patch.object(module, "_map_tasks", commit_before_dispatch): + with patch.object(RaySearchTaskExecutor, "map", commit_before_dispatch): assert query.to_list(execution="ray", concurrency=2) == expected assert committed == [True] assert query.to_list(execution="ray") == query.to_list() diff --git a/paimon-python/pypaimon/tests/ray_hybrid_parallel_test.py b/paimon-python/pypaimon/tests/ray_hybrid_parallel_test.py index 755e6a2e53c7..9c08c8e6a5f7 100644 --- a/paimon-python/pypaimon/tests/ray_hybrid_parallel_test.py +++ b/paimon-python/pypaimon/tests/ray_hybrid_parallel_test.py @@ -24,15 +24,27 @@ ray = pytest.importorskip("ray") -from pypaimon.ray import full_text_search, vector_search +from pypaimon.execution.python.read.search import hybrid_search +from pypaimon.ray import full_text_search +from pypaimon.ray.search_execution import RaySearchTaskExecutor +from pypaimon.table.source.search_execution import SearchTask -def builder(count): - routes = [SimpleNamespace(route=SimpleNamespace(is_vector=lambda: False), search_builder=i) - for i in range(count)] - return SimpleNamespace(route_builders=lambda: routes, - to_route_result=lambda route, result: (route.search_builder, result), - rank=lambda results: results) +def builder(count, execute): + def search_builder(ordinal): + return SimpleNamespace( + ordinal=ordinal, + new_full_text_scan=lambda: SimpleNamespace(scan=lambda: ordinal), + new_full_text_read=lambda: SimpleNamespace(read_plan=execute)) + + routes = [SimpleNamespace(route=SimpleNamespace(is_vector=lambda: False), + search_builder=search_builder(i)) for i in range(count)] + result = SimpleNamespace( + route_builders=lambda: routes, + to_route_result=lambda route, result: (route.search_builder.ordinal, result), + rank=lambda results: results) + result.execute = lambda executor=None: hybrid_search.execute(result, executor) + return result @pytest.mark.parametrize("concurrency, count", [(2, 2), (3, 2), (4, 2), (2, 4)]) @@ -44,9 +56,10 @@ def test_routes_overlap_within_budget_and_preserve_order(concurrency, count): peak = 0 completed = [] - def execute(ordinal, *, concurrency, ray_remote_args): + def execute(ordinal, *, executor): nonlocal active, peak - assert ray_remote_args == {"num_cpus": 0.5} + assert executor._remote_args == {"num_cpus": 0.5} + concurrency = executor.concurrency with lock: active += concurrency peak = max(peak, active) @@ -60,9 +73,8 @@ def execute(ordinal, *, concurrency, ray_remote_args): second_completed.set() return ordinal - with patch.object(full_text_search, "_execute_full_text_search", execute): - result = full_text_search._execute_hybrid_search( - builder(count), concurrency=concurrency, ray_remote_args={"num_cpus": 0.5}) + result = full_text_search._execute_hybrid_search( + builder(count, execute), concurrency=concurrency, ray_remote_args={"num_cpus": 0.5}) assert result == [(i, i) for i in range(count)] assert completed.index(1) < completed.index(0) assert 2 <= peak <= concurrency @@ -72,14 +84,13 @@ def execute(ordinal, *, concurrency, ray_remote_args): def test_serial_routes_keep_full_budget(concurrency, count): calls = [] - def execute(ordinal, **kwargs): - assert kwargs["concurrency"] == concurrency + def execute(ordinal, *, executor): + assert executor.concurrency == concurrency calls.append(ordinal) return ordinal - with patch.object(full_text_search, "_execute_full_text_search", execute): - assert full_text_search._execute_hybrid_search(builder(count), concurrency=concurrency) == [ - (i, i) for i in range(count)] + assert full_text_search._execute_hybrid_search( + builder(count, execute), concurrency=concurrency) == [(i, i) for i in range(count)] assert calls == list(range(count)) @@ -87,7 +98,7 @@ def test_route_failure_cancels_sibling_ray_tasks(): waiting = Event() submitted, cancelled = [], [] - def submit(context, item): + def submit(operation, context, item, *dependencies): submitted.append(item) return item @@ -101,32 +112,32 @@ def wait(refs, **kwargs): Event().wait(0.01) return [], refs - def execute(ordinal, **kwargs): + def execute(ordinal, *, executor): if ordinal == 0: assert waiting.wait(10) raise ValueError("route failed") - return list(vector_search._map_tasks(None, None, [0, 1, 2], 1, {})) + return list(executor.map([SearchTask(None, None, item) for item in range(3)])) - with patch.object(full_text_search, "_execute_full_text_search", execute), \ - patch.object(ray, "put", return_value=None), \ + with patch.object(ray, "put", return_value=None), \ patch.object(ray, "remote", return_value=remote), \ patch.object(ray, "wait", side_effect=wait), \ patch.object(ray, "cancel", side_effect=cancelled.append), \ pytest.raises(ValueError, match="route failed"): - full_text_search._execute_hybrid_search(builder(2), concurrency=2) + full_text_search._execute_hybrid_search(builder(2, execute), concurrency=2) assert submitted == [0] assert cancelled == [0] def test_cancellation_scope_restores_state_and_skips_pending_routes(): + executor = RaySearchTaskExecutor() cancelled = Event() cancelled.set() - with pytest.raises(CancelledError), vector_search._task_cancellation(cancelled): + with pytest.raises(CancelledError), executor.cancellation_scope(cancelled): pytest.fail("cancelled route started") - assert getattr(vector_search._task_state, "cancelled", None) is None + assert executor._cancellations.cancellation_events() == () outer, inner = Event(), Event() - with vector_search._task_cancellation(outer): - with vector_search._task_cancellation(inner): - assert vector_search._task_state.cancelled is inner - assert vector_search._task_state.cancelled is outer - assert vector_search._task_state.cancelled is None + with executor.cancellation_scope(outer): + with executor.cancellation_scope(inner): + assert executor._cancellations.cancellation_events() == (outer, inner) + assert executor._cancellations.cancellation_events() == (outer,) + assert executor._cancellations.cancellation_events() == () diff --git a/paimon-python/pypaimon/tests/ray_search_execution_test.py b/paimon-python/pypaimon/tests/ray_search_execution_test.py new file mode 100644 index 000000000000..ec337c55067a --- /dev/null +++ b/paimon-python/pypaimon/tests/ray_search_execution_test.py @@ -0,0 +1,345 @@ +# 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. + +"""Ray transports opaque tasks and shares the query's global task budget.""" + +import subprocess +import sys +from concurrent.futures import CancelledError, ThreadPoolExecutor +from contextlib import closing +from threading import Event, Lock +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +ray = pytest.importorskip("ray") + +from pypaimon.ray.search_execution import RaySearchTaskExecutor +from pypaimon.table.source.search_execution import SearchTask + + +def test_ray_facades_do_not_load_python_or_native_execution(): + source = """ +import importlib.abc +import sys +from types import SimpleNamespace + +class RejectEngines(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.startswith(('pypaimon.execution.python', 'pypaimon.execution.native', 'pypaimon_rust')): + raise AssertionError('Ray facade selected an execution engine: ' + fullname) + +sys.meta_path.insert(0, RejectEngines()) +from pypaimon.ray.vector_search import _execute_vector_search +from pypaimon.ray.batch_vector_search import _execute_batch_vector_search +from pypaimon.ray.full_text_search import _execute_full_text_search, _execute_hybrid_search +from pypaimon.ray.search_execution import RaySearchTaskExecutor, _execute_task + +plan = object() +result = object() +def read(requested, executor): + assert requested is plan + assert isinstance(executor, RaySearchTaskExecutor) + return result +reader = SimpleNamespace(read_plan=read, read_batch_plan=read) +scan = SimpleNamespace(scan=lambda: plan) +builder = SimpleNamespace(new_vector_search_read=lambda: reader, + new_batch_vector_search_read=lambda: reader, + new_full_text_read=lambda: reader, + new_vector_search_scan=lambda: scan, + new_full_text_scan=lambda: scan, + execute=lambda executor: read(plan, executor)) +for execute in (_execute_vector_search, _execute_batch_vector_search, + _execute_full_text_search, _execute_hybrid_search): + assert execute(builder, concurrency=2) is result +assert _execute_task(lambda context, item, dependency: context + item + dependency, + b'context', b'item', b'dependency') == b'contextitemdependency' +""" + result = subprocess.run([sys.executable, '-c', source], stdout=subprocess.PIPE, + stderr=subprocess.PIPE, universal_newlines=True, timeout=30) + assert result.returncode == 0, result.stdout + result.stderr + + +def test_shared_dependencies_resolve_once_without_dependent_worker_allocation(): + contexts, submitted, values = [], [], {} + context = object() + + def put(value): + contexts.append(value) + return ('context', len(contexts)) + + def submit(operation, context_ref, item, *dependencies): + ref = len(submitted) + submitted.append((operation, context_ref, item, dependencies)) + return ref + + def get(ref): + if ref not in values: + operation, context_ref, item, dependencies = submitted[ref] + value = contexts[context_ref[1] - 1] + values[ref] = operation(value, item, *(get(dependency) for dependency in dependencies)) + return values[ref] + + remote = SimpleNamespace(remote=submit) + remote.options = lambda **kwargs: remote + dependency = SearchTask(lambda shared, item: item, context, b'scalar-result') + work = [SearchTask(lambda shared, item, scalar: item + scalar, context, item, (dependency,)) + for item in (b'first-', b'second-')] + with patch.object(ray, 'put', side_effect=put), \ + patch.object(ray, 'remote', return_value=remote), \ + patch.object(ray, 'wait', side_effect=lambda refs, **kwargs: (refs[:1], refs[1:])), \ + patch.object(ray, 'get', side_effect=get), patch.object(ray, 'cancel'): + assert list(RaySearchTaskExecutor(2).map(work, ordered=True)) == [ + (0, b'first-scalar-result'), (1, b'second-scalar-result')] + assert contexts == [context] + assert len(submitted) == 3 + # The dependency is a top-level argument so Ray can resolve it before + # allocating either dependent worker, even if each task needs all CPUs. + assert submitted[1][3] == submitted[2][3] == (0,) + + +def test_concurrent_callback_maps_share_global_budget_and_release_slots_on_close(): + executor = RaySearchTaskExecutor(2) + child = executor.with_concurrency(2) + first_wait, budget_wait, release = Event(), Event(), Event() + lock = Lock() + submitted, outstanding, peak = [], set(), 0 + + def submit(operation, context, item, *dependencies): + nonlocal peak + with lock: + ref = len(submitted) + submitted.append(item) + outstanding.add(ref) + peak = max(peak, len(outstanding)) + return ref + + def wait(refs, **kwargs): + first_wait.set() + assert release.wait(10) + return refs[:1], refs[1:] + + def get(ref): + with lock: + outstanding.remove(ref) + return ref + + original_wait = executor._budget.wait + + def no_slots(): + budget_wait.set() + original_wait() + + def run(view, prefix): + tasks = [SearchTask(None, None, (prefix, item)) for item in range(4)] + with closing(view.map(tasks)) as results: + return list(results) + + remote = SimpleNamespace(remote=submit) + remote.options = lambda **kwargs: remote + with patch.object(ray, 'put', return_value=None), \ + patch.object(ray, 'remote', return_value=remote), \ + patch.object(ray, 'wait', side_effect=wait), \ + patch.object(ray, 'get', side_effect=get), patch.object(ray, 'cancel'), \ + patch.object(executor._budget, 'wait', side_effect=no_slots), \ + ThreadPoolExecutor(max_workers=2) as pool: + left = pool.submit(run, executor, 'left') + assert first_wait.wait(10) + right = pool.submit(run, child, 'right') + try: + assert budget_wait.wait(10) + assert submitted == [('left', 0), ('left', 1)] + finally: + release.set() + assert len(left.result(timeout=10)) == len(right.result(timeout=10)) == 4 + assert peak == 2 + # A closed map must release unconsumed pending and buffered slots. + with closing(child.map([SearchTask(None, None, item) for item in range(4)])) as results: + next(results) + outstanding.clear() + assert len(run(executor, 'reused')) == 4 + + +def test_cancellation_scope_reaches_native_callbacks_on_another_thread(): + executor = RaySearchTaskExecutor(2) + cancelled, waiting = Event(), Event() + remote = SimpleNamespace(remote=lambda *args: 'pending') + remote.options = lambda **kwargs: remote + + def wait(refs, **kwargs): + assert kwargs['timeout'] == 0.1 + waiting.set() + cancelled.wait(10) + return [], refs + + def run(): + return list(child.map([SearchTask(None, None, 'item')])) + + with patch.object(ray, 'put', return_value=None), \ + patch.object(ray, 'remote', return_value=remote), \ + patch.object(ray, 'wait', side_effect=wait), patch.object(ray, 'cancel') as cancel, \ + executor.cancellation_scope(cancelled), ThreadPoolExecutor(max_workers=1) as pool: + child = executor.with_concurrency(1) + future = pool.submit(run) + try: + assert waiting.wait(10) + finally: + cancelled.set() + with pytest.raises(CancelledError): + future.result(timeout=10) + cancel.assert_called_once_with('pending') + + +@pytest.mark.parametrize('family', ['vector', 'full-text']) +def test_python_primary_key_readers_reject_unsupported_executor_explicitly(family): + 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.read_plan(object(), executor=RaySearchTaskExecutor(2)) + + +class _CallbackDriver: + def __init__(self): + self.items = [] + self.outstanding = set() + self.peak = 0 + self.lock = Lock() + + def submit(self, operation, context, item, *dependencies): + with self.lock: + ref = len(self.items) + self.items.append(item) + self.outstanding.add(ref) + self.peak = max(self.peak, len(self.outstanding)) + return ref + + def get(self, ref): + with self.lock: + self.outstanding.remove(ref) + return self.items[ref] + + def cancel(self, ref): + with self.lock: + self.outstanding.discard(ref) + + def remote(self): + remote = SimpleNamespace(remote=self.submit) + remote.options = lambda **kwargs: remote + return remote + + +def test_concurrent_reader_queries_do_not_share_cancellation_but_share_budget(): + from pypaimon.execution.python.read.search.distributed_search import create_distributed_reader + from pypaimon.execution.python.read.search.vector_search_read import DataEvolutionVectorRead + + base = RaySearchTaskExecutor(2) + table = SimpleNamespace(options=SimpleNamespace(data_evolution_enabled=lambda: True)) + reader = DataEvolutionVectorRead(table, 1, SimpleNamespace(name='embedding'), [1.]) + events = {name: Event() for name in ('cancelled', 'completed')} + waiting = {name: Event() for name in events} + release = Event() + driver = _CallbackDriver() + + def wait(refs, **kwargs): + name = driver.items[refs[0]] + waiting[name].set() + if name == 'cancelled': + assert events[name].wait(10) + return [], refs + assert release.wait(10) + return refs[:1], refs[1:] + + def run(name): + # This is the query view selected by the public Pure read_plan path. + view = create_distributed_reader(reader, base)._executor + with view.cancellation_scope(events[name]): + return list(view.map([SearchTask(None, None, name)])) + + with patch.object(ray, 'put', return_value=None), \ + patch.object(ray, 'remote', return_value=driver.remote()), \ + patch.object(ray, 'wait', side_effect=wait), \ + patch.object(ray, 'get', side_effect=driver.get), \ + patch.object(ray, 'cancel', side_effect=driver.cancel), ThreadPoolExecutor(max_workers=2) as pool: + left = pool.submit(run, 'cancelled') + try: + assert waiting['cancelled'].wait(10) + right = pool.submit(run, 'completed') + assert waiting['completed'].wait(10) + events['cancelled'].set() + with pytest.raises(CancelledError): + left.result(timeout=10) + assert not right.done() + finally: + events['cancelled'].set() + release.set() + assert right.result(timeout=10) == [(0, 'completed')] + assert driver.peak == base.concurrency + assert base._cancellations.cancellation_events() == () + + +def test_hybrid_failure_does_not_cancel_another_query_using_same_executor(): + from pypaimon.execution.python.read.search import hybrid_search + from pypaimon.tests.ray_hybrid_parallel_test import builder + + base = RaySearchTaskExecutor(2) + waiting = {name: Event() for name in ('failed', 'completed')} + release = Event() + driver = _CallbackDriver() + + def wait(refs, **kwargs): + name, ordinal = driver.items[refs[0]] + waiting[name].set() + if name == 'failed': + Event().wait(0.01) + return [], refs + assert release.wait(10) + return refs[:1], refs[1:] + + def run_route(name, ordinal, executor): + if name == 'failed' and ordinal == 0: + assert waiting['failed'].wait(10) + assert waiting['completed'].wait(10) + raise ValueError('first query failed') + return list(executor.map([SearchTask(None, None, (name, ordinal))])) + + def run(name): + query = builder(2, lambda ordinal, executor: run_route(name, ordinal, executor)) + return hybrid_search.execute(query, base) + + with patch.object(ray, 'put', return_value=None), \ + patch.object(ray, 'remote', return_value=driver.remote()), \ + patch.object(ray, 'wait', side_effect=wait), \ + patch.object(ray, 'get', side_effect=driver.get), \ + patch.object(ray, 'cancel', side_effect=driver.cancel), ThreadPoolExecutor(max_workers=2) as pool: + left = pool.submit(run, 'failed') + try: + assert waiting['failed'].wait(10) + right = pool.submit(run, 'completed') + assert waiting['completed'].wait(10) + with pytest.raises(ValueError, match='first query failed'): + left.result(timeout=10) + assert not right.done() + finally: + release.set() + assert right.result(timeout=10) == [ + (ordinal, [(0, ('completed', ordinal))]) for ordinal in range(2)] + assert driver.peak == base.concurrency + assert base._cancellations.cancellation_events() == () diff --git a/paimon-python/pypaimon/tests/ray_vector_filter_exactness_test.py b/paimon-python/pypaimon/tests/ray_vector_filter_exactness_test.py index 95dddc57d123..ed1d3c02cb1f 100644 --- a/paimon-python/pypaimon/tests/ray_vector_filter_exactness_test.py +++ b/paimon-python/pypaimon/tests/ray_vector_filter_exactness_test.py @@ -15,12 +15,16 @@ # specific language governing permissions and limitations # under the License. +from dataclasses import replace from unittest.mock import patch import pytest pytest.importorskip("ray") +from pypaimon.ray.search_execution import RaySearchTaskExecutor +from pypaimon.tests.search_execution_test_helpers import task_parts + from pypaimon.read.table_read import TableRead from pypaimon.execution.python.read.search.vector_search_read import AbstractVectorSearchReadImpl from pypaimon.tests import ray_vector_search_test as ray_fixtures @@ -38,14 +42,14 @@ def test_ray_applies_exact_row_filter_before_worker_top_k(table, ray_cluster, ba scalar_index(table) table.raw_table = table.raw_table.copy({ "vector-index.search-mode": mode, "global-index.filter.refine-from-data": str(refine).lower()}) - from pypaimon.ray import batch_vector_search, vector_search from pypaimon.execution.python.read.search import global_index_live_row_filter - module = batch_vector_search if batch else vector_search - original_map = module._map_tasks + original_map = RaySearchTaskExecutor.map calls = [] - def record(worker, context, items, *args): + def record(executor, tasks, ordered=False): + worker, context, items = task_parts(tasks) + def checked(context, split, scalar_filter=None): # These patches run in the actual Ray worker, not on the driver. original = AbstractVectorSearchReadImpl._matching_candidate_rows @@ -73,11 +77,13 @@ def verify(reader, candidates, snapshot): result = worker(context, split, scalar_filter) return result, observed - for ordinal, (result, observed) in original_map(checked, context, items, *args): + mapped = original_map(executor, [replace(task, operation=checked) for task in tasks], + ordered=ordered) + for ordinal, (result, observed) in mapped: calls.extend(observed) yield ordinal, result - with patch.object(module, "_map_tasks", record), \ + with patch.object(RaySearchTaskExecutor, "map", record), \ patch.object(AbstractVectorSearchReadImpl, "_pre_filters", side_effect=AssertionError("driver must not build filters")), \ patch.object(global_index_live_row_filter, "live_rows", @@ -95,7 +101,6 @@ def verify(reader, candidates, snapshot): def test_worker_filters_overlapping_scalar_shards_at_planned_snapshot(tmp_path, ray_cluster, batch, kind): import pyarrow as pa import pypaimon.multimodal as pm - from pypaimon.ray import batch_vector_search, vector_search pytest.importorskip("paimon_vindex") schema = pa.schema([("id", pa.int64()), ("name", pa.string()), @@ -120,17 +125,16 @@ def test_worker_filters_overlapping_scalar_shards_at_planned_snapshot(tmp_path, expected = [{"id": 0}, {"id": 4}, {"id": 6}] expected = [expected, expected] if batch else expected assert search.to_list() == expected - module = batch_vector_search if batch else vector_search - original_map = module._map_tasks + original_map = RaySearchTaskExecutor.map dispatched = [] - def commit_before_workers(worker, context, splits, *args): + def commit_before_workers(executor, tasks, ordered=False): if not dispatched: table.delete("id = 0") dispatched.append(True) - return original_map(worker, context, splits, *args) + return original_map(executor, tasks, ordered=ordered) - with patch.object(module, "_map_tasks", commit_before_workers), \ + with patch.object(RaySearchTaskExecutor, "map", commit_before_workers), \ patch.object(AbstractVectorSearchReadImpl, "_pre_filters", side_effect=AssertionError("driver filter")): result = search.to_list(execution="ray", concurrency=2) @@ -148,7 +152,7 @@ def test_shared_scalar_index_is_evaluated_once_on_ray(table, tmp_path, ray_clust from pathlib import Path import pyarrow as pa - from pypaimon.ray import vector_search + from pypaimon.execution.python.read.search import distributed_vector_search schema = pa.schema([("id", pa.int64()), ("name", pa.string()), ("embedding", pa.list_(pa.float32(), 2))]) @@ -162,14 +166,14 @@ def test_shared_scalar_index_is_evaluated_once_on_ray(table, tmp_path, ray_clust options["btree-index.fallback-scan-max-size"] = "0 b" table.raw_table = table.raw_table.copy(options) marker = str(tmp_path / "scalar-evaluations") - original = vector_search._prepare_scalar_filter + original = distributed_vector_search._prepare_scalar_filter def record(context, split): with open(marker, "a") as output: output.write(str(os.getpid()) + "\n") return original(context, split) - with patch.object(vector_search, "_prepare_scalar_filter", record), \ + with patch.object(distributed_vector_search, "_prepare_scalar_filter", record), \ patch.object(AbstractVectorSearchReadImpl, "_scalar_index_result", side_effect=AssertionError("driver scalar evaluation")): # Each task requires the whole two-CPU cluster. Pending searches must diff --git a/paimon-python/pypaimon/tests/ray_vector_refinement_test.py b/paimon-python/pypaimon/tests/ray_vector_refinement_test.py index 386dbaafbc92..f984df8fbc22 100644 --- a/paimon-python/pypaimon/tests/ray_vector_refinement_test.py +++ b/paimon-python/pypaimon/tests/ray_vector_refinement_test.py @@ -21,7 +21,10 @@ pytest.importorskip("ray") -from pypaimon.ray import vector_search as search_module +from pypaimon.ray.search_execution import RaySearchTaskExecutor +from pypaimon.tests.search_execution_test_helpers import task_parts + +from pypaimon.execution.python.read.search import distributed_vector_search as search_module from pypaimon.read.table_read import TableRead from pypaimon.tests import ray_vector_search_test as fixtures from pypaimon.tests.ray_vector_search_test import VECTORS, add_rows, build_index, ids @@ -42,20 +45,21 @@ def test_refinement_reads_vectors_only_on_workers(table, ray_cluster, metric, co query = table.search([1., 1.], options={"refine_factor": "2"}).select(["id"]).limit(3) expected = ids(query.to_arrow()) original_read = TableRead._new_arrow_batch_reader - original_map = search_module._map_tasks + original_map = RaySearchTaskExecutor.map dispatched = [] def no_driver_vectors(read, *args, **kwargs): assert "embedding" not in [field.name for field in read.read_type] return original_read(read, *args, **kwargs) - def record(worker, context, items, *args): + def record(executor, tasks, ordered=False): + worker, context, items = task_parts(tasks) if worker is search_module._search_raw_split: dispatched.append(len(items)) - return original_map(worker, context, items, *args) + return original_map(executor, tasks, ordered=ordered) with patch.object(TableRead, "_new_arrow_batch_reader", no_driver_vectors), \ - patch.object(search_module, "_map_tasks", record): + patch.object(RaySearchTaskExecutor, "map", record): assert ids(query.to_arrow(execution="ray", concurrency=concurrency)) == expected assert dispatched == [2] @@ -65,7 +69,7 @@ def test_candidate_intersection_excludes_nan_and_streams_real_splits(table, ray_ add_rows(table, [[float("nan"), 1.], [1., 1.]], 2) query = table.search([1., 1.], options={"metric": "l2"}).limit(1)._for_execution() reader = query._vector_search_builder(query).new_vector_search_read() - distributed = search_module._RayVectorSearchRead(reader, 2, {}) + distributed = search_module.DistributedVectorSearchRead(reader, RaySearchTaskExecutor(2, {})) candidates = RoaringBitmap64() candidates.add(0) candidates.add(3) @@ -84,9 +88,9 @@ def test_candidate_intersection_excludes_nan_and_streams_real_splits(table, ray_ def test_empty_refinement_does_not_plan_or_dispatch(table): query = table.search([1., 1.]).limit(1)._for_execution() reader = query._vector_search_builder(query).new_vector_search_read() - distributed = search_module._RayVectorSearchRead(reader, 2, {}) + distributed = search_module.DistributedVectorSearchRead(reader, RaySearchTaskExecutor(2, {})) with patch.object(distributed, "_plan_raw_read", side_effect=AssertionError("empty read")), \ - patch.object(search_module, "_map_tasks", side_effect=AssertionError("empty task")): + patch.object(RaySearchTaskExecutor, "map", side_effect=AssertionError("empty task")): result = distributed._read_raw_search( [Range(0, 100)], None, [1., 1.], score_candidates=RoaringBitmap64()) assert result.results().is_empty() @@ -118,16 +122,17 @@ def test_commit_between_candidate_selection_and_refinement_keeps_snapshot(table, table.raw_table = table.raw_table.copy({"global-index.column-update-action": "DROP_PARTITION_INDEX"}) query = table.search([1., 1.], options={"refine_factor": "2"}).select(["id"]).limit(1) expected = ids(query.to_arrow()) - original = search_module._map_tasks + original = RaySearchTaskExecutor.map changed = [] - def commit_before_refine(worker, *args): + def commit_before_refine(executor, tasks, ordered=False): + worker, context, items = task_parts(tasks) if worker is search_module._search_raw_split and not changed: table.update("id = 0", {"embedding": [100., 1.]}) changed.append(True) - return original(worker, *args) + return original(executor, tasks, ordered=ordered) - with patch.object(search_module, "_map_tasks", commit_before_refine): + with patch.object(RaySearchTaskExecutor, "map", commit_before_refine): assert ids(query.to_arrow(execution="ray")) == expected assert changed == [True] assert ids(query.to_arrow()) != expected diff --git a/paimon-python/pypaimon/tests/ray_vector_search_test.py b/paimon-python/pypaimon/tests/ray_vector_search_test.py index a90b00c8205b..b6e1c74f5ca0 100644 --- a/paimon-python/pypaimon/tests/ray_vector_search_test.py +++ b/paimon-python/pypaimon/tests/ray_vector_search_test.py @@ -26,9 +26,12 @@ ray = pytest.importorskip("ray") +from pypaimon.ray.search_execution import RaySearchTaskExecutor +from pypaimon.tests.search_execution_test_helpers import execute_tasks, task_parts + import pypaimon.multimodal as pm from pypaimon.multimodal.query import ScanQuery -from pypaimon.ray import vector_search as search_module +from pypaimon.execution.python.read.search import distributed_vector_search as search_module from pypaimon.snapshot.time_travel_util import TimeTravelUtil from pypaimon.execution.python.read.search.vector_search_scan import DataEvolutionVectorScan @@ -92,14 +95,15 @@ def test_raw_search_matches_exact_top_k(table, ray_cluster, metric, concurrency) scores = vectors.dot(query_vector) expected = sorted(sorted(range(len(VECTORS)), key=lambda i: (-scores[i], i))[:3]) assert ids(query.to_arrow()) == expected - original = search_module._map_tasks + original = RaySearchTaskExecutor.map dispatched = [] - def record(worker, context, items, *args): + def record(executor, tasks, ordered=False): + worker, context, items = task_parts(tasks) dispatched.append(len(items)) - return original(worker, context, items, *args) + return original(executor, tasks, ordered=ordered) - with patch.object(search_module, "_map_tasks", record): + with patch.object(RaySearchTaskExecutor, "map", record): assert ids(query.to_arrow(execution="ray", concurrency=concurrency)) == expected assert dispatched == [2] @@ -117,13 +121,14 @@ def test_indexed_search_and_raw_fallback_use_persisted_metric(table, ray_cluster query = table.search([1., 1.], options=options).select(["id"]).limit(3) expected = ids(query.to_arrow()) observed = [] - original = search_module._map_tasks + original = RaySearchTaskExecutor.map - def record(worker, context, items, *args): + def record(executor, tasks, ordered=False): + worker, context, items = task_parts(tasks) observed.append((worker.__name__, len(items))) - return original(worker, context, items, *args) + return original(executor, tasks, ordered=ordered) - with patch.object(search_module, "_map_tasks", record): + with patch.object(RaySearchTaskExecutor, "map", record): assert ids(query.to_arrow(execution="ray", concurrency=2)) == expected assert ("_search_index_split", 2) in observed assert any(name == "_search_raw_split" and count > 0 for name, count in observed) @@ -174,19 +179,19 @@ def test_query_pins_worker_reads_and_lookup_but_can_be_reused(table, ray_cluster add_rows(table, [[1., 1.], [10., 1.]]) query = table.search([1., 1.]).select(["id", "category"]).limit(1) expected = query.to_arrow().to_pylist() - original = search_module._map_tasks + original = RaySearchTaskExecutor.map changed = [] - def change_before_dispatch(*args): + def change_before_dispatch(executor, tasks, ordered=False): if not changed: if change == "delete": table.delete("id = 0") else: table.update("id = 0", {"category": "changed"}) changed.append(True) - return original(*args) + return original(executor, tasks, ordered=ordered) - with patch.object(search_module, "_map_tasks", change_before_dispatch): + with patch.object(RaySearchTaskExecutor, "map", change_before_dispatch): assert query.to_arrow(execution="ray").to_pylist() == expected latest = query.to_arrow().to_pylist() assert latest != expected @@ -202,7 +207,7 @@ def append_before_scan(search_scan): return scan(search_scan) with patch.object(DataEvolutionVectorScan, "scan", append_before_scan), \ - patch.object(search_module, "_map_tasks", side_effect=AssertionError("empty task")): + patch.object(RaySearchTaskExecutor, "map", side_effect=AssertionError("empty task")): result = query.to_arrow(execution="ray") assert result.schema.names == ["id"] assert result.num_rows == 0 @@ -248,7 +253,7 @@ def inspect_context(table, ordinal): time.sleep(0.2) return os.getpid(), TimeTravelUtil.resolve_snapshot(table).id - results = list(search_module._map_tasks(inspect_context, pinned, [0, 1], 2, {"num_cpus": 1})) + results = list(execute_tasks(inspect_context, pinned, [0, 1], 2, {"num_cpus": 1})) assert len({result[0] for _, result in results}) == 2 assert all(pid != os.getpid() and snapshot == 2 for _, (pid, snapshot) in results) assert ids(query.to_arrow(execution="ray")) == [0, 1] @@ -271,7 +276,7 @@ def retry_once(table, marker): os.close(fd) raise RuntimeError("retry this read") - assert list(search_module._map_tasks( + assert list(execute_tasks( retry_once, pinned, [marker], 1, {"max_retries": 1, "retry_exceptions": True}, ordered=ordered)) == [(0, (1, [0]))] @@ -286,7 +291,7 @@ def fail(context, item): with patch.object(ray, "cancel", wraps=ray.cancel) as cancel: with pytest.raises(ValueError, match="injected search failure"): - list(search_module._map_tasks(fail, None, [0, 1, 2], 2, {}, ordered=ordered)) + list(execute_tasks(fail, None, [0, 1, 2], 2, {}, ordered=ordered)) assert cancel.call_count == 2 @@ -303,7 +308,7 @@ def record_wait(refs, **kwargs): return wait(refs, **kwargs) with patch.object(ray, "wait", record_wait), patch.object(ray, "cancel", wraps=ray.cancel) as cancel: - with closing(search_module._map_tasks(echo, None, list(range(10)), 2, {}, ordered=ordered)) as tasks: + with closing(execute_tasks(echo, None, list(range(10)), 2, {}, ordered=ordered)) as tasks: next(tasks) assert counts == [2] or (ordered and counts == [2, 1]) assert cancel.call_count == 2 - len(counts) @@ -316,22 +321,23 @@ def test_refinement_uses_global_candidates(table): add_rows(table, [[0., 1.], [10., 1.], [20., 1.], [1., 1.]]) query = table.search([1., 1.], options={"refine_factor": "2"}).limit(1)._for_execution() reader = query._vector_search_builder(query).new_vector_search_read() - distributed = search_module._RayVectorSearchRead(reader, 2, {}) + distributed = search_module.DistributedVectorSearchRead(reader, RaySearchTaskExecutor(2, {})) index_file = IndexFileMeta("ivf-flat", "fake", 1, 2, None) splits = [IndexVectorSearchSplit(0, 1, [index_file]), IndexVectorSearchSplit(2, 3, [index_file])] - def completed(worker, context, items, *args): + def completed(executor, tasks, ordered=False): + worker, context, items = task_parts(tasks) if worker is search_module._search_raw_split: for ordinal, split in enumerate(items): yield ordinal, worker(context, split) return # Row 3 is the true nearest, but is outside the GLOBAL approximate top-2. # Refining independently per shard would incorrectly bring it back in. - assert args[2] is True # Index tasks request plan-ordered results. + assert ordered is True # Index tasks request plan-ordered results. yield 0, ("l2", {0: 10., 1: 9.}) yield 1, ("l2", {2: 8., 3: 7.}) - with patch.object(search_module, "_map_tasks", completed): + with patch.object(RaySearchTaskExecutor, "map", completed): result = distributed._read_indexed(splits, [1., 1.], query._table._read_snapshot) assert list(result.results()) == [0] assert result.score_getter()(0) == 0.5 # L2 score is 1 / (1 + squared distance). @@ -343,16 +349,16 @@ def test_duplicate_scores_keep_plan_order(table): query = table.search([1., 1.]).limit(2)._for_execution() reader = query._vector_search_builder(query).new_vector_search_read() - distributed = search_module._RayVectorSearchRead(reader, 2, {}) + distributed = search_module.DistributedVectorSearchRead(reader, RaySearchTaskExecutor(2, {})) index_file = IndexFileMeta("ivf-flat", "fake", 1, 2) splits = [IndexVectorSearchSplit(0, 1, [index_file]), IndexVectorSearchSplit(0, 3, [index_file])] - def completed(*args): - assert args[5] is True + def completed(executor, tasks, ordered=False): + assert ordered is True yield 0, ("l2", {0: 1., 1: 2.}) yield 1, ("l2", {0: 100.}) - with patch.object(search_module, "_map_tasks", completed): + with patch.object(RaySearchTaskExecutor, "map", completed): result = distributed._read_indexed(splits, [1., 1.], None) assert result.score_getter()(0) == 1. assert list(result.top_k(1).results()) == [1] @@ -363,17 +369,17 @@ def test_metric_mismatch_fails_even_when_one_shard_has_no_hits(table): query = table.search([1., 1.]) reader = query._vector_search_builder(query).new_vector_search_read() - distributed = search_module._RayVectorSearchRead(reader, 2, {}) + distributed = search_module.DistributedVectorSearchRead(reader, RaySearchTaskExecutor(2, {})) closed = [] - def completed(*args): + def completed(executor, tasks, ordered=False): try: yield 0, ("l2", {}) yield 1, ("inner_product", {1: 10.}) finally: closed.append(True) - with patch.object(search_module, "_map_tasks", completed), \ + with patch.object(RaySearchTaskExecutor, "map", completed), \ pytest.raises(ValueError, match="different metrics"): list(distributed._search_index_splits( [IndexVectorSearchSplit(0, 1, []), IndexVectorSearchSplit(2, 3, [])], [1., 1.], 1, None)) @@ -419,7 +425,7 @@ def test_nan_scores_fail_before_worker_top_k(table, ray_cluster): def test_slow_first_task_bounds_completed_results(ordered, early_close): submitted, completed, consumed, cancelled = [], [], [], [] - def submit(context, item): + def submit(operation, context, item, *dependencies): # A slow prefix must not allow all later tasks/results to accumulate. assert len(submitted) - len(consumed) < 3 submitted.append(item) @@ -438,7 +444,7 @@ def wait(refs, **kwargs): patch.object(ray, "wait", side_effect=wait), \ patch.object(ray, "get", side_effect=lambda ref: ref), \ patch.object(ray, "cancel", side_effect=cancelled.append): - with closing(search_module._map_tasks(None, None, list(range(10)), 3, {}, ordered)) as tasks: + with closing(execute_tasks(None, None, list(range(10)), 3, {}, ordered)) as tasks: for ordinal, result in tasks: assert ordinal == result consumed.append(ordinal) diff --git a/paimon-python/pypaimon/tests/search_execution_test_helpers.py b/paimon-python/pypaimon/tests/search_execution_test_helpers.py new file mode 100644 index 000000000000..c185bda6f981 --- /dev/null +++ b/paimon-python/pypaimon/tests/search_execution_test_helpers.py @@ -0,0 +1,33 @@ +# 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. + +"""Build and inspect public search tasks in focused execution tests.""" + +from pypaimon.ray.search_execution import RaySearchTaskExecutor +from pypaimon.table.source.search_execution import SearchTask + + +def task_parts(tasks): + tasks = list(tasks) + first = tasks[0] + assert all(task.operation is first.operation and task.context is first.context for task in tasks) + return first.operation, first.context, [task.item for task in tasks] + + +def execute_tasks(operation, context, items, concurrency, remote_args, ordered=False): + executor = RaySearchTaskExecutor(concurrency, remote_args) + return executor.map([SearchTask(operation, context, item) for item in items], ordered=ordered) diff --git a/paimon-python/pypaimon/tests/search_task_execution_test.py b/paimon-python/pypaimon/tests/search_task_execution_test.py new file mode 100644 index 000000000000..bc90e00de2d9 --- /dev/null +++ b/paimon-python/pypaimon/tests/search_task_execution_test.py @@ -0,0 +1,219 @@ +# 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. + +"""Search scheduling preserves backend selection and opaque task ownership.""" + +from concurrent.futures import CancelledError +from concurrent.futures import ThreadPoolExecutor +from threading import Event +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest + +from pypaimon.execution.native.read.native_search_plan import NativeSearchRead +from pypaimon.execution.native.read.search_execution import native_task_executor +from pypaimon.table.source.search_dispatch import _BatchSearchRead, _SplitSearchRead +from pypaimon.table.source.search_execution import SearchTaskExecutor +from pypaimon.table.source.hybrid_search_builder import HybridSearchBuilderImpl + + +class InlineExecutor(SearchTaskExecutor): + def __init__(self): + self.tasks = [] + self.closed = 0 + + def map(self, tasks, ordered=False): + assert ordered + self.tasks.extend(tasks) + try: + for ordinal, task in enumerate(tasks): + yield ordinal, task.execute() + finally: + self.closed += 1 + + +class MapExecutor(SearchTaskExecutor): + def __init__(self, map_): + self._map = map_ + + def map(self, tasks, ordered=False): + return self._map(tasks, ordered) + + +def test_hybrid_execute_without_scheduler_retains_a_local_instance_override(): + result = object() + builder = HybridSearchBuilderImpl(None) + builder.execute_local = Mock(return_value=result) + assert builder.execute() is result + builder.execute_local.assert_called_once_with() + + +@pytest.mark.parametrize('family', ['single', 'batch', 'full-text']) +def test_native_search_schedules_opaque_tasks_without_constructing_python_reader(family): + table = SimpleNamespace(options=SimpleNamespace( + options=SimpleNamespace(to_map=lambda: {}), data_evolution_enabled=lambda: True)) + plan = object() + executor = InlineExecutor() + result = SimpleNamespace(row_ids=lambda: {7: 0.75}) + + def read(native_plan, executor): + assert native_plan is plan + assert list(executor(b'shared-context', [b'index-task', b'raw-task'])) == [b'index-result', b'raw-result'] + return [result, result] if family == 'batch' else result + + core_read = SimpleNamespace(read_plan=read, read_batch_plan=read) + native = NativeSearchRead(table, core_read, False, ()) + native.resolve_plan = lambda requested: plan + python_factory = Mock(side_effect=AssertionError('Python executor was constructed')) + reader = (_BatchSearchRead if family == 'batch' else _SplitSearchRead)(native, python_factory) + + def execute_task(payload, shared): + assert shared == b'shared-context' + return payload.replace(b'task', b'result') + + core_table = SimpleNamespace(_execute_search_task=execute_task) + with patch('pypaimon.execution.native.read.search_execution._native_table', return_value=core_table): + results = (reader.read_batch_plan(plan, executor=executor) if family == 'batch' + else [reader.read_plan(plan, executor=executor)]) + assert [{row: value.score_getter()(row) for row in value.results()} for value in results] == ( + [{7: 0.75}, {7: 0.75}] if family == 'batch' else [{7: 0.75}]) + assert [task.item for task in executor.tasks] == [b'index-task', b'raw-task'] + assert all(task.context is executor.tasks[0].context for task in executor.tasks) + assert executor.tasks[0].context == (table, b'shared-context') + assert executor.closed == 1 + python_factory.assert_not_called() + + +@pytest.mark.parametrize('values', [[], [(1, b'wrong-order')], [(0, b'ok'), (0, b'duplicate')]]) +def test_native_task_adapter_rejects_missing_or_misordered_results(values): + executor = MapExecutor(lambda tasks, ordered: iter(values)) + with pytest.raises(ValueError, match='Search task executor'): + list(native_task_executor(object(), executor)(b'context', [b'one-task'])) + + +def test_scheduler_cancellation_is_not_replaced_by_cleanup_failure_or_python_replay(): + failure = CancelledError('query cancelled') + + class FailedResults: + def __iter__(self): + return self + + def __next__(self): + raise failure + + def close(self): + raise RuntimeError('cleanup failed') + + executor = MapExecutor(lambda tasks, ordered: FailedResults()) + with pytest.raises(CancelledError) as caught: + list(native_task_executor(object(), executor)(b'context', [b'task'])) + assert caught.value is failure + + +def test_distributed_native_failure_does_not_start_python_search(): + failure = CancelledError('remote task failed') + native = SimpleNamespace(resolve_plan=lambda plan: plan, read_plan=Mock(side_effect=failure)) + fallback = Mock(side_effect=AssertionError('Python search replayed')) + reader = _SplitSearchRead(native, fallback) + with pytest.raises(CancelledError) as caught: + reader.read_plan(object(), executor=InlineExecutor()) + assert caught.value is failure + fallback.assert_not_called() + + +def test_core_failure_prevents_additional_work_from_its_scheduler_callback(): + table = SimpleNamespace(options=SimpleNamespace( + options=SimpleNamespace(to_map=lambda: {}), data_evolution_enabled=lambda: True)) + failure = ValueError('inconsistent index metrics') + callbacks = [] + + def read(plan, executor): + callbacks.append(executor) + raise failure + + reader = NativeSearchRead(table, SimpleNamespace(read_plan=read), False, ()) + executor = InlineExecutor() + with pytest.raises(ValueError) as caught: + reader.read_plan(object(), executor=executor) + assert caught.value is failure + with pytest.raises(CancelledError): + list(callbacks[0](b'context', [b'additional-task'])) + # A sibling callback can finish after Core has already failed and cancelled + # the query. Its synthetic cancellation cannot replace the Core failure. + callbacks[0].raise_if_failed() + assert executor.tasks == [] + + +def test_sibling_cancellation_during_slow_cleanup_preserves_the_original_worker_error(): + failure = ValueError('worker could not read the pinned index') + closing, release = Event(), Event() + + class FailedResults: + def __iter__(self): + return self + + def __next__(self): + raise failure + + def close(self): + closing.set() + assert release.wait(10) + + def read(plan, executor): + with ThreadPoolExecutor(max_workers=1) as pool: + failed = pool.submit(lambda: list(executor(b'context', [b'failed-task']))) + assert closing.wait(10) + try: + # Model the SDK receiving the sibling cancellation first while + # the actual failure is still closing its pending resources. + list(executor(b'context', [b'sibling-task'])) + finally: + release.set() + with pytest.raises(ValueError) as original: + failed.result(timeout=10) + assert original.value is failure + + table = SimpleNamespace(options=SimpleNamespace( + options=SimpleNamespace(to_map=lambda: {}), data_evolution_enabled=lambda: True)) + reader = NativeSearchRead(table, SimpleNamespace(read_plan=read), False, ()) + with pytest.raises(ValueError) as caught: + reader.read_plan(object(), executor=MapExecutor(lambda tasks, ordered: FailedResults())) + assert caught.value is failure + + +def test_closing_callback_iterator_does_not_replace_a_core_error_with_generator_exit(): + callback = native_task_executor(object(), MapExecutor(lambda tasks, ordered: iter([(0, b'result')]))) + results = callback(b'context', [b'task']) + assert next(results) == b'result' + results.close() + callback.raise_if_failed() + + +def test_closing_callback_after_a_core_error_does_not_promote_a_cleanup_failure(): + class FailedCleanup: + def __iter__(self): + return iter([(0, b'result')]) + + def close(self): + raise RuntimeError('worker cleanup failed') + + callback = native_task_executor(object(), MapExecutor(lambda tasks, ordered: FailedCleanup())) + results = callback(b'context', [b'task']) + assert next(results) == b'result' + results.close() + callback.raise_if_failed() From 76304dc400a40f371958e4452d6e2c3a8af174b9 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Wed, 21 Oct 2026 21:13:03 +0800 Subject: [PATCH 2/9] Keep source plans on the selected search execution backend --- docs/docs/pypaimon/multimodal-search.md | 16 ++ paimon-python/pypaimon/execution/README.md | 8 + .../native/read/native_hybrid_search.py | 2 +- .../native/read/native_search_plan.py | 219 +++++++++------ .../read/search/primary_key_full_text_scan.py | 2 +- .../read/search/primary_key_vector_scan.py | 2 +- .../pypaimon/table/source/search_dispatch.py | 32 +-- .../pypaimon/table/source/search_types.py | 12 +- .../tests/lazy_search_execution_test.py | 256 ++++++++++++++---- .../pypaimon/tests/native_search_plan_test.py | 66 +++-- .../tests/search_task_execution_test.py | 20 +- 11 files changed, 442 insertions(+), 193 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-search.md b/docs/docs/pypaimon/multimodal-search.md index 9daa8153148f..30544083259e 100644 --- a/docs/docs/pypaimon/multimodal-search.md +++ b/docs/docs/pypaimon/multimodal-search.md @@ -359,6 +359,22 @@ results = builder.new_batch_vector_search_read().read_batch_plan( plan, executor=RaySearchTaskExecutor(concurrency=4)) ``` +Plans contain ordinary snapshot and split metadata. A copied, pickled, or +Python-constructed plan can be passed to a Native reader without choosing the +Python backend. For example, to execute an explicitly selected set of sources: + +```python +from pypaimon.table.source.search_types import VectorSearchScanPlan + +selected = VectorSearchScanPlan(plan.splits()[:1], plan.snapshot()) +results = builder.new_batch_vector_search_read().read_batch_plan(selected) +``` + +The supplied splits are authoritative: readers do not add uncovered raw ranges +or rediscover other search indexes. Query vectors, limits and filters belong to +the reader. Rust Core coordinates global candidate selection and refinement; +its internal shard protocol is separate from the public source plan. + Hybrid builders expose `execute(executor=...)` for the same scheduler. The existing `execute_local()` methods retain their local defaults. diff --git a/paimon-python/pypaimon/execution/README.md b/paimon-python/pypaimon/execution/README.md index 532eb5c5038e..5ae09dd3136d 100644 --- a/paimon-python/pypaimon/execution/README.md +++ b/paimon-python/pypaimon/execution/README.md @@ -154,6 +154,14 @@ opaque shared context and task bytes; Core owns both their interpretation and global query coordination. The global candidate limit is applied before refinement, and full-text raw fallback retains one corpus for BM25 statistics. +Vector and full-text plans are snapshot-and-split data in both backends. Rust +plans have public constructors and can be serialized; the Native adapter builds +them from ordinary Python plans, including plans made by Python scans or sent +to another process. Readers own the table, schema, query and filter. Native +selection is retained after transport, and source subsets never add missing +indexes or raw ranges. Shard callbacks use an internal binding bridge; they are +not part of Rust's public Plan/Read API. + ## Corresponding execution APIs Compare the operation's inputs, results and lifecycle when changing a backend. diff --git a/paimon-python/pypaimon/execution/native/read/native_hybrid_search.py b/paimon-python/pypaimon/execution/native/read/native_hybrid_search.py index d7c485fcfac3..d8ac3578299c 100644 --- a/paimon-python/pypaimon/execution/native/read/native_hybrid_search.py +++ b/paimon-python/pypaimon/execution/native/read/native_hybrid_search.py @@ -71,7 +71,7 @@ def try_native_hybrid_search(builder, executor=None): from pypaimon.execution.native.read.search_execution import native_task_executor callback = native_task_executor(table, executor) try: - result = native.execute(executor=callback) + result = native._execute_with_executor(callback) finally: callback.cancel() callback.raise_if_failed() diff --git a/paimon-python/pypaimon/execution/native/read/native_search_plan.py b/paimon-python/pypaimon/execution/native/read/native_search_plan.py index be49b2660197..a470957f696c 100644 --- a/paimon-python/pypaimon/execution/native/read/native_search_plan.py +++ b/paimon-python/pypaimon/execution/native/read/native_search_plan.py @@ -15,17 +15,14 @@ # specific language governing permissions and limitations # under the License. -"""Native search plans with ordinary, transportable Python source metadata.""" +"""Translate ordinary snapshot-and-split plans to and from Rust Core.""" -import hashlib import json -import pickle +from copy import copy, deepcopy from pypaimon.common.json_util import JSON -from pypaimon.common.options.core_options import CoreOptions -from pypaimon.common.options.options import Options from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult -from pypaimon.read.split_serializer import deserialize_split_v1 +from pypaimon.read.split_serializer import deserialize_split_v1, serialize_split_v1 from pypaimon.execution.native.table import (_partition_fields) from pypaimon.snapshot.snapshot import Snapshot from pypaimon.table.source.search_types import ( @@ -35,44 +32,7 @@ ) from pypaimon.table.source.primary_key_scored_result import PrimaryKeyScoredResult from pypaimon.utils.range import Range -from pypaimon.write.commit_message_serializer import _index_from_row - - -def _metadata_fingerprint(plan): - # Only captured metadata participates: never pickle a Table, catalog or - # live native handle. This also detects edits through a shallow copy. - metadata = (plan.snapshot(), plan.splits(), getattr(plan, 'snapshot_id', None)) - return hashlib.sha256(pickle.dumps(metadata, protocol=4)).digest() - - -class _NativePlan: - def __getstate__(self): - # Transport immutable snapshot data and ordinary Python splits. The - # original result and its metadata remain unchanged by serialization. - state = self.__dict__.copy() - state.pop('_native_plan', None) - state.pop('_native_fingerprint', None) - return state - - def _unchanged_native_plan(self): - native = getattr(self, '_native_plan', None) - return native if native is not None and self._native_fingerprint == _metadata_fingerprint(self) else None - - -class _NativeVectorPlan(_NativePlan, VectorSearchScanPlan): - pass - - -class _NativeFullTextPlan(_NativePlan, FullTextScanPlan): - pass - - -class _NativePkVectorPlan(_NativePlan, PrimaryKeyVectorScanPlan): - pass - - -class _NativePkFullTextPlan(_NativePlan, PrimaryKeyFullTextScanPlan): - pass +from pypaimon.write.commit_message_serializer import _index_from_row, _index_row def _range(value): @@ -83,11 +43,11 @@ def _files(values): return [_index_from_row(bytes(value)) for value in values] -def _plan_from_native(table, native, full_text, primary_key, source_modes): +def _plan_from_native(table, native, full_text, primary_key): from pypaimon.table.source.full_text_search_split import IndexFullTextSearchSplit, RawFullTextSearchSplit from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit, RawVectorSearchSplit - metadata = json.loads(native._metadata()) + metadata = json.loads(native.metadata()) snapshot = JSON.from_json(json.dumps(metadata['snapshot']), Snapshot) if metadata['snapshot'] is not None else None splits = [] for value in metadata['splits']: @@ -108,19 +68,67 @@ def _plan_from_native(table, native, full_text, primary_key, source_modes): splits.append(RawFullTextSearchSplit(ranges) if full_text else RawVectorSearchSplit( ranges, _files(value['scalar_index_files']), value['index_type'])) if primary_key: - cls = _NativePkFullTextPlan if full_text else _NativePkVectorPlan - plan = cls(metadata['snapshot_id'] or 0, splits) - plan._snapshot = snapshot + cls = PrimaryKeyFullTextScanPlan if full_text else PrimaryKeyVectorScanPlan + plan = cls(metadata['snapshot_id'] or 0, splits, snapshot) else: - cls = _NativeFullTextPlan if full_text else _NativeVectorPlan + cls = FullTextScanPlan if full_text else VectorSearchScanPlan plan = cls(splits, snapshot) - plan._native_plan = native - plan._source_modes = source_modes - plan._table_options = dict(table.options.options.to_map()) - plan._native_fingerprint = _metadata_fingerprint(plan) return plan +def _split_metadata(split, full_text, column): + from pypaimon.table.source.full_text_search_split import IndexFullTextSearchSplit, RawFullTextSearchSplit + from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit, RawVectorSearchSplit + + expected = ((PrimaryKeyFullTextSearchSplit, IndexFullTextSearchSplit, RawFullTextSearchSplit) if full_text else + (PrimaryKeyVectorSearchSplit, IndexVectorSearchSplit, RawVectorSearchSplit)) + if not isinstance(split, expected): + raise TypeError('Unsupported %s search split: %s' % ( + 'full-text' if full_text else 'vector', type(split).__name__)) + + def files(values): + return [list(_index_row(value)) for value in values] + + def range_(value): + return {'from': value.from_, 'to': value.to} + + if isinstance(split, (PrimaryKeyVectorSearchSplit, PrimaryKeyFullTextSearchSplit)): + return {'kind': 'primary-key', 'data_split': list(serialize_split_v1(split.data_split)), + 'payloads': files(split.payloads), 'uncovered_data_files': list(split.uncovered_data_files), + 'row_ranges_by_file': {name: [range_(value) for value in values] + for name, values in split.row_ranges_by_file.items()}} + if isinstance(split, (IndexVectorSearchSplit, IndexFullTextSearchSplit)): + return {'kind': 'index', 'column': split.column_name if full_text else column, + 'row_range': {'from': split.row_range_start, 'to': split.row_range_end}, + 'index_files': files(split.full_text_index_files if full_text else split.vector_index_files), + 'scalar_index_files': files(split.scalar_index_files)} + if isinstance(split, RawVectorSearchSplit): + return {'kind': 'raw', 'row_ranges': [range_(value) for value in split.row_ranges], + 'scalar_index_files': files(split.scalar_index_files), 'index_type': split.index_type} + if isinstance(split, RawFullTextSearchSplit): + return {'kind': 'raw', 'row_ranges': [range_(value) for value in split.row_ranges], + 'scalar_index_files': [], 'index_type': None} + raise TypeError('Unsupported search split: %s' % type(split).__name__) + + +def _plan_to_native(plan, full_text, column=None): + from pypaimon_rust.datafusion import FullTextScanPlan as NativeFullTextPlan, VectorScanPlan as NativeVectorPlan + + snapshot = plan.snapshot() + return (NativeFullTextPlan if full_text else NativeVectorPlan)( + None if snapshot is None else JSON.to_json(snapshot), + json.dumps([_split_metadata(split, full_text, column) for split in plan.splits()])) + + +def _capture_builder(builder): + # Public builders contain only table configuration. Freeze query values and + # options for reader transport without copying the catalog or live SDK cache. + captured = copy(builder) + captured.__dict__ = {key: value if key == '_table' else deepcopy(value) + for key, value in builder.__dict__.items()} + return captured + + def _result(table, value, primary_key): return (PrimaryKeyScoredResult._from_native(table, value) if primary_key else DictBasedScoredIndexResult(value.row_ids())) @@ -129,12 +137,13 @@ def _result(table, value, primary_key): class NativeSearchScan: """Convert Rust scan metadata without selecting another executor.""" - def __init__(self, table, scan, full_text, primary_key, source_modes): + def __init__(self, table, scan, full_text, primary_key, builder=None, batch=False): self._table = table self._scan = scan self._full_text = full_text self._primary_key = primary_key - self._source_modes = source_modes + self._builder = builder + self._batch = batch def __getstate__(self): state = self.__dict__.copy() @@ -142,50 +151,77 @@ def __getstate__(self): return state def available(self): - return self._scan is not None + return self._scan is not None or self._builder is not None def scan(self): - return _plan_from_native(self._table, self._scan.scan(), self._full_text, self._primary_key, self._source_modes) + if self._scan is None: + native = _builder_to_native(self._builder, self._full_text, self._batch, scan=True) + self._scan = native.new_full_text_scan() if self._full_text else native.new_vector_search_scan() + return _plan_from_native(self._table, self._scan.scan(), self._full_text, self._primary_key) class NativeSearchRead: - """Read a compatible live Rust plan; transport fallback belongs to dispatch.""" + """Convert source metadata and delegate complete queries to Rust Core.""" - def __init__(self, table, native, primary_key, source_modes): + def __init__(self, table, native, primary_key, full_text=False, builder=None, batch=False): self._table = table self._native = native self._primary_key = primary_key - self._source_modes = source_modes - self._table_options = dict(table.options.options.to_map()) + self._full_text = full_text + self._builder = builder + self._batch = batch def __getstate__(self): state = self.__dict__.copy() state['_native'] = None return state + def _get_native_reader(self): + if self._native is None: + native = _builder_to_native(self._builder, self._full_text, self._batch) + self._native = (native.new_full_text_read() if self._full_text else + native.new_batch_vector_search_read() if self._batch else + native.new_vector_search_read()) + return self._native + + def _plan(self, plan): + if (self._primary_key and isinstance(plan, (PrimaryKeyVectorScanPlan, PrimaryKeyFullTextScanPlan)) + and plan.snapshot() is None and plan.snapshot_id): + snapshot = self._table.snapshot_manager().get_snapshot_by_id(plan.snapshot_id) + if snapshot is None: + raise ValueError('Captured search snapshot %s is unavailable' % plan.snapshot_id) + plan = copy(plan) + plan._snapshot = snapshot + column = None if self._builder is None or self._full_text else self._builder._vector_column.name + return _plan_to_native(plan, self._full_text, column) + def read_plan(self, plan, executor=None): + native = self._get_native_reader() + plan = self._plan(plan) if executor is None: - result = self._native.read_plan(plan) + result = native.read_plan(plan) else: self._validate_task_execution() from pypaimon.execution.native.read.search_execution import native_task_executor callback = native_task_executor(self._table, executor) try: - result = self._native.read_plan(plan, executor=callback) + result = native._read_plan_with_executor(plan, callback) finally: callback.cancel() callback.raise_if_failed() return _result(self._table, result, self._primary_key) def read_batch_plan(self, plan, executor=None): + native = self._get_native_reader() + plan = self._plan(plan) if executor is None: - results = self._native.read_batch_plan(plan) + results = native.read_batch_plan(plan) else: self._validate_task_execution() from pypaimon.execution.native.read.search_execution import native_task_executor callback = native_task_executor(self._table, executor) try: - results = self._native.read_batch_plan(plan, executor=callback) + results = native._read_batch_plan_with_executor(plan, callback) finally: callback.cancel() callback.raise_if_failed() @@ -196,22 +232,36 @@ def _validate_task_execution(self): if self._primary_key or not self._table.options.data_evolution_enabled(): raise ValueError('Distributed search supports only data-evolution tables') - def resolve_plan(self, plan): - # Search modes choose source work at scan time. A new reader may use - # another query or limit, but must still consume these captured splits. - if (self._native is not None and isinstance(plan, _NativePlan) - and plan._source_modes == self._source_modes - and plan._table_options == self._table_options): - return plan._unchanged_native_plan() - return None + def read(self, splits): + return self.read_plan(self._split_plan(splits)) + + def read_batch(self, splits): + return self.read_batch_plan(self._split_plan(splits)) + def _split_plan(self, splits): + from pypaimon.snapshot.time_travel_util import TimeTravelUtil + snapshot = TimeTravelUtil.resolve_snapshot(self._table) + if self._primary_key: + cls = PrimaryKeyFullTextScanPlan if self._full_text else PrimaryKeyVectorScanPlan + return cls(0 if snapshot is None else snapshot.id, splits, snapshot) + return (FullTextScanPlan if self._full_text else VectorSearchScanPlan)(splits, snapshot) -def _source_modes(builder, full_text=False): - options = dict(builder._table.options.options.to_map()) - options.update(getattr(builder, '_options', {})) - core = CoreOptions(Options(options)) - return ((core.full_text_index_search_mode(),) if full_text else - (core.vector_index_search_mode(), core.scalar_index_search_mode())) + +def _builder_to_native(builder, full_text, batch=False, scan=False): + if full_text: + from pypaimon.execution.native.read.native_full_text_search import _native_full_text_builder + native = _native_full_text_builder(builder, require_limit=not scan) + else: + from pypaimon.execution.native.read.native_vector_search import _native_vector_builder + native = _native_vector_builder(builder, batch) + if native is not None and not scan: + if batch: + native.with_query_vectors(builder._query_vectors) + else: + native.with_query_vector(builder._query_vector) + if native is None: + raise RuntimeError('Native search runtime is unavailable for the captured reader configuration') + return native def native_vector_scan(builder, batch=False): @@ -219,7 +269,7 @@ def native_vector_scan(builder, batch=False): native = _native_vector_builder(builder, batch) return None if native is None else NativeSearchScan( builder._table, native.new_vector_search_scan(), False, - builder._is_primary_key_vector_search(), _source_modes(builder)) + builder._is_primary_key_vector_search(), _capture_builder(builder), batch) def native_vector_read(builder, batch=False): @@ -234,7 +284,8 @@ def native_vector_read(builder, batch=False): native.with_query_vector(builder._query_vector) reader = native.new_vector_search_read() return NativeSearchRead( - builder._table, reader, builder._is_primary_key_vector_search(), _source_modes(builder)) + builder._table, reader, builder._is_primary_key_vector_search(), False, + _capture_builder(builder), batch) def native_full_text_scan(builder): @@ -242,7 +293,7 @@ def native_full_text_scan(builder): native = _native_full_text_builder(builder, require_limit=False) return None if native is None else NativeSearchScan( builder._table, native.new_full_text_scan(), True, - builder._primary_key_full_text_definition() is not None, _source_modes(builder, True)) + builder._primary_key_full_text_definition() is not None, _capture_builder(builder)) def native_full_text_read(builder): @@ -250,4 +301,4 @@ def native_full_text_read(builder): native = _native_full_text_builder(builder) return None if native is None else NativeSearchRead( builder._table, native.new_full_text_read(), - builder._primary_key_full_text_definition() is not None, _source_modes(builder, True)) + builder._primary_key_full_text_definition() is not None, True, _capture_builder(builder)) diff --git a/paimon-python/pypaimon/execution/python/read/search/primary_key_full_text_scan.py b/paimon-python/pypaimon/execution/python/read/search/primary_key_full_text_scan.py index 290b5e3bee30..b66e29c90283 100644 --- a/paimon-python/pypaimon/execution/python/read/search/primary_key_full_text_scan.py +++ b/paimon-python/pypaimon/execution/python/read/search/primary_key_full_text_scan.py @@ -61,7 +61,7 @@ def scan(self): entry, self._definition.field_id, self._definition.index_type, self._partition_filter)) return PrimaryKeyFullTextScanPlan( - snapshot.id, _full_text_bucket_splits(source_splits, entries)) + snapshot.id, _full_text_bucket_splits(source_splits, entries), snapshot) def _matches(entry, field_id, index_type, partition_filter): diff --git a/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_scan.py b/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_scan.py index d2ca03291f1b..7d2c0ea334f5 100644 --- a/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_scan.py +++ b/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_scan.py @@ -87,7 +87,7 @@ def scan(self): lambda entry: _matches(entry, self._vector_column.id, index_type, self._partition_filter)) return PrimaryKeyVectorScanPlan( - snapshot.id, _bucket_splits(source_splits, entries)) + snapshot.id, _bucket_splits(source_splits, entries), snapshot) def _matches(entry, field_id, index_type, partition_filter): diff --git a/paimon-python/pypaimon/table/source/search_dispatch.py b/paimon-python/pypaimon/table/source/search_dispatch.py index 14446063941d..b17445026494 100644 --- a/paimon-python/pypaimon/table/source/search_dispatch.py +++ b/paimon-python/pypaimon/table/source/search_dispatch.py @@ -17,12 +17,10 @@ """Coordinate search execution and transport without coupling the backends. -A live compatible plan executes in Rust. Transported plans and explicit split -reads select Python lazily. An execution error never triggers another search. +The selected backend consumes ordinary snapshot-and-split plans, including +transported plans. An execution error never triggers another search. """ -from copy import copy - class _SearchScan: def __init__(self, native, python_factory): @@ -49,37 +47,23 @@ def _python_reader(self): self._python_executor = self._python_factory() return self._python_executor - def _reader_for_plan(self, plan): - reader = copy(self._python_reader()) - if plan.snapshot() is not None: - reader._table = reader._table._copy_with_snapshot(plan.snapshot()) - return reader - def read_plan(self, plan, executor=None): - native = self._native_reader.resolve_plan(plan) - if native is not None: - return (self._native_reader.read_plan(native) if executor is None else - self._native_reader.read_plan(native, executor=executor)) - reader = self._reader_for_plan(plan) - return reader.read_plan(plan) if executor is None else reader.read_plan(plan, executor=executor) + return (self._native_reader.read_plan(plan) if executor is None else + self._native_reader.read_plan(plan, executor=executor)) def read_batch_plan(self, plan, executor=None): - native = self._native_reader.resolve_plan(plan) - if native is not None: - return (self._native_reader.read_batch_plan(native) if executor is None else - self._native_reader.read_batch_plan(native, executor=executor)) - reader = self._reader_for_plan(plan) - return reader.read_batch_plan(plan) if executor is None else reader.read_batch_plan(plan, executor=executor) + return (self._native_reader.read_batch_plan(plan) if executor is None else + self._native_reader.read_batch_plan(plan, executor=executor)) class _SplitSearchRead(_SearchRead): def read(self, splits): - return self._python_reader().read(splits) + return self._native_reader.read(splits) class _BatchSearchRead(_SearchRead): def read_batch(self, splits): - return self._python_reader().read_batch(splits) + return self._native_reader.read_batch(splits) def python_search_reader(reader): diff --git a/paimon-python/pypaimon/table/source/search_types.py b/paimon-python/pypaimon/table/source/search_types.py index a744afa518c5..5e846f67ee9f 100644 --- a/paimon-python/pypaimon/table/source/search_types.py +++ b/paimon-python/pypaimon/table/source/search_types.py @@ -141,8 +141,10 @@ class PrimaryKeyVectorSearchSplit: class PrimaryKeyVectorScanPlan(VectorSearchScanPlan): - def __init__(self, snapshot_id, splits): - super().__init__(splits) + def __init__(self, snapshot_id, splits, snapshot=None): + if snapshot is not None and snapshot.id != snapshot_id: + raise ValueError('Primary-key plan snapshot ID does not match its captured snapshot') + super().__init__(splits, snapshot) self.snapshot_id = snapshot_id @@ -155,6 +157,8 @@ class PrimaryKeyFullTextSearchSplit: class PrimaryKeyFullTextScanPlan(FullTextScanPlan): - def __init__(self, snapshot_id, splits): - super().__init__(splits) + def __init__(self, snapshot_id, splits, snapshot=None): + if snapshot is not None and snapshot.id != snapshot_id: + raise ValueError('Primary-key plan snapshot ID does not match its captured snapshot') + super().__init__(splits, snapshot) self.snapshot_id = snapshot_id diff --git a/paimon-python/pypaimon/tests/lazy_search_execution_test.py b/paimon-python/pypaimon/tests/lazy_search_execution_test.py index 6923756a3b6b..1eeb3830f480 100644 --- a/paimon-python/pypaimon/tests/lazy_search_execution_test.py +++ b/paimon-python/pypaimon/tests/lazy_search_execution_test.py @@ -19,16 +19,16 @@ from contextlib import ExitStack import importlib +import json import pickle -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from unittest.mock import Mock, patch import pyarrow as pa import pytest from pypaimon import CatalogFactory, Schema -from pypaimon.table.source import search_dispatch -from pypaimon.table.source.search_types import VectorSearchScanPlan +from pypaimon.table.source.search_types import FullTextScanPlan, VectorSearchScanPlan @pytest.fixture @@ -60,7 +60,7 @@ def _factories(builder, family): def _native_builder(family): - plan = SimpleNamespace(_metadata=lambda: '{"snapshot":null,"snapshot_id":null,"splits":[]}') + plan = SimpleNamespace(metadata=lambda: '{"snapshot":null,"snapshot_id":null,"splits":[]}') result = SimpleNamespace(row_ids=lambda: {7: 0.5}) reader = SimpleNamespace(read_plan=Mock(return_value=result), read_batch_plan=Mock(return_value=[result, result])) scan = SimpleNamespace(scan=Mock(return_value=plan)) @@ -92,21 +92,41 @@ def _forbid_python_execution(): return stack +def _ordinary_plan(family): + return (FullTextScanPlan if family == 'text' else VectorSearchScanPlan)([], None) + + @pytest.mark.parametrize('family', ['vector', 'batch', 'text']) -def test_native_search_prepares_scans_reads_and_serializes_without_python(search_table, family): +def test_native_search_prepares_plain_plans_and_transports_readers_without_python(search_table, family): builder = _builder(search_table, family) native, core_scan, core_read = _native_builder(family) scan_factory, read_factory = _factories(builder, family) - with _forbid_python_execution(), patch(_native_factory(family), return_value=native): - scan, read = scan_factory(), read_factory() + native_plan = object() + with _forbid_python_execution(), patch(_native_factory(family), return_value=native), patch( + 'pypaimon.execution.native.read.native_search_plan._plan_to_native', return_value=native_plan) as convert: + scan, reader = scan_factory(), read_factory() plan = scan.scan() - results = read.read_batch_plan(plan) if family == 'batch' else [read.read_plan(plan)] + assert type(plan) is (FullTextScanPlan if family == 'text' else VectorSearchScanPlan) + assert not hasattr(plan, '_native_plan') + results = reader.read_batch_plan(plan) if family == 'batch' else [reader.read_plan(plan)] assert [list(result.results()) for result in results] == [[7]] * (2 if family == 'batch' else 1) - # Dropping live handles for transport must not create fallback readers. - assert pickle.loads(pickle.dumps(read))._native_reader._native is None - assert not pickle.loads(pickle.dumps(scan))._native_scan.available() - core_scan.scan.assert_called_once_with() - (core_read.read_batch_plan if family == 'batch' else core_read.read_plan).assert_called_once_with(plan._native_plan) + # The transported reader reconstructs only the selected Native backend. + transported = pickle.loads(pickle.dumps(reader)) + assert transported._native_reader._native is None + copied_plan = pickle.loads(pickle.dumps(plan)) + copied_results = (transported.read_batch_plan(copied_plan) if family == 'batch' + else [transported.read_plan(copied_plan)]) + assert [list(result.results()) for result in copied_results] == [[7]] * (2 if family == 'batch' else 1) + assert transported._python_executor is None + copied_scan = pickle.loads(pickle.dumps(scan)) + assert copied_scan._native_scan.available() + assert copied_scan.scan().splits() == [] + assert copied_scan._python_executor is None + assert convert.call_count == 2 + assert convert.call_args[0][0].splits() == [] + assert core_scan.scan.call_args_list == [((),), ((),)] + method = core_read.read_batch_plan if family == 'batch' else core_read.read_plan + assert method.call_args_list == [((native_plan,),), ((native_plan,),)] @pytest.mark.parametrize('family', ['vector', 'batch', 'text']) @@ -127,25 +147,29 @@ def test_native_search_reader_introspection_does_not_select_python_execution(sea @pytest.mark.parametrize('family', ['vector', 'batch', 'text']) -def test_native_search_split_read_selects_python_only_when_invoked(search_table, family): +def test_native_search_split_read_delegates_to_native_without_python(search_table, family): builder = _builder(search_table, family) - native, _, _ = _native_builder(family) + native, _, core_read = _native_builder(family) _, read_factory = _factories(builder, family) + native_plan = object() method_name = 'read_batch' if family == 'batch' else 'read' - python_method = Mock(return_value='split fallback') - python = SimpleNamespace(**{method_name: python_method}) - factory = ('_new_python_full_text_read' if family == 'text' else - '_new_python_batch_vector_read' if family == 'batch' else '_new_python_vector_read') - with patch(_native_factory(family), return_value=native), \ - patch.object(type(builder), factory, return_value=python) as prepare: + with _forbid_python_execution(), patch(_native_factory(family), return_value=native), patch( + 'pypaimon.execution.native.read.native_search_plan._plan_to_native', + return_value=native_plan) as convert, patch( + 'pypaimon.snapshot.time_travel_util.TimeTravelUtil.resolve_snapshot', + return_value=None) as snapshot: reader = read_factory() split_read = getattr(reader, method_name) - prepare.assert_not_called() - splits = ['transported split'] - assert split_read(splits) == 'split fallback' - assert split_read(splits) == 'split fallback' - prepare.assert_called_once_with() - assert python_method.call_args_list == [((splits,),), ((splits,),)] + splits = [] + for _ in range(2): + results = split_read(splits) if family == 'batch' else [split_read(splits)] + assert [list(result.results()) for result in results] == [[7]] * (2 if family == 'batch' else 1) + assert reader._python_executor is None + assert snapshot.call_count == 2 + assert convert.call_count == 2 + assert all(call[0][0].splits() is splits and call[0][0].snapshot() is None for call in convert.call_args_list) + method = core_read.read_batch_plan if family == 'batch' else core_read.read_plan + assert method.call_args_list == [((native_plan,),), ((native_plan,),)] @pytest.mark.parametrize('phase', ['scan', 'read']) @@ -153,7 +177,8 @@ def test_native_search_failure_never_creates_python_execution(search_table, phas builder = _builder(search_table, 'vector') native, core_scan, core_read = _native_builder('vector') failure = OSError('search may already have read remote state') - with _forbid_python_execution(), patch(_native_factory('vector'), return_value=native): + with _forbid_python_execution(), patch(_native_factory('vector'), return_value=native), patch( + 'pypaimon.execution.native.read.native_search_plan._plan_to_native', return_value=object()): scan, reader = builder.new_vector_search_scan(), builder.new_vector_search_read() if phase == 'scan': core_scan.scan.side_effect = failure @@ -167,32 +192,39 @@ def test_native_search_failure_never_creates_python_execution(search_table, phas assert caught.value is failure -def test_python_plan_fallback_is_created_once_when_actually_needed(search_table): - builder = _builder(search_table, 'vector') - native, _, _ = _native_builder('vector') - python = SimpleNamespace(_table=search_table, read_plan=Mock(return_value='fallback')) - with patch(_native_factory('vector'), return_value=native), \ - patch.object(type(builder), '_new_python_vector_read', return_value=python) as prepare: - reader = builder.new_vector_search_read() - prepare.assert_not_called() - plan = VectorSearchScanPlan([], None) - assert reader.read_plan(plan) == 'fallback' - assert reader.read_plan(plan) == 'fallback' - prepare.assert_called_once_with() - assert python.read_plan.call_count == 2 +@pytest.mark.parametrize('family', ['vector', 'batch', 'text']) +def test_ordinary_python_plans_use_native_without_fallback(search_table, family): + builder = _builder(search_table, family) + native, _, core_read = _native_builder(family) + _, read_factory = _factories(builder, family) + native_plan = object() + with _forbid_python_execution(), patch(_native_factory(family), return_value=native), patch( + 'pypaimon.execution.native.read.native_search_plan._plan_to_native', return_value=native_plan) as convert: + reader = read_factory() + plan = _ordinary_plan(family) + results = reader.read_batch_plan(plan) if family == 'batch' else [reader.read_plan(plan)] + assert [list(result.results()) for result in results] == [[7]] * (2 if family == 'batch' else 1) + assert reader._python_executor is None + assert convert.call_args[0][0] is plan + method = core_read.read_batch_plan if family == 'batch' else core_read.read_plan + method.assert_called_once_with(native_plan) -def test_native_reader_keeps_the_query_options_captured_at_construction(search_table): +def test_native_reader_keeps_query_options_captured_before_transport(search_table): builder = _builder(search_table, 'vector').with_option('ivf-flat.metric', 'l2') native, _, _ = _native_builder('vector') with patch(_native_factory('vector'), return_value=native): reader = builder.new_vector_search_read() builder.with_option('ivf-flat.metric', 'cosine') - fallback = search_dispatch.python_search_reader(reader) - assert fallback._options == {'ivf-flat.metric': 'l2'} + transported = pickle.loads(pickle.dumps(reader)) + with _forbid_python_execution(), patch(_native_factory('vector'), return_value=native) as prepare, patch( + 'pypaimon.execution.native.read.native_search_plan._plan_to_native', return_value=object()): + transported.read_plan(_ordinary_plan('vector')) + assert prepare.call_args[0][0]._options == {'ivf-flat.metric': 'l2'} + assert transported._python_executor is None -def test_batch_reader_freezes_query_sequence_before_lazy_fallback(search_table): +def test_batch_reader_freezes_query_sequence_before_transport(search_table): queries = [[1.0, 0.0]] builder = _builder(search_table, 'batch').with_query_vectors(queries) native, _, _ = _native_builder('batch') @@ -201,8 +233,27 @@ def test_batch_reader_freezes_query_sequence_before_lazy_fallback(search_table): queries.append([0.0, 1.0]) queries[0][:] = [3.0, 4.0] assert native.with_query_vectors.call_args[0][0] == [[1.0, 0.0]] - fallback = search_dispatch.python_search_reader(pickle.loads(pickle.dumps(reader))) - assert fallback._query_vectors == [[1.0, 0.0]] + transported = pickle.loads(pickle.dumps(reader)) + with _forbid_python_execution(), patch(_native_factory('batch'), return_value=native), patch( + 'pypaimon.execution.native.read.native_search_plan._plan_to_native', return_value=object()): + transported.read_batch_plan(_ordinary_plan('batch')) + assert native.with_query_vectors.call_args[0][0] == [[1.0, 0.0]] + assert transported._python_executor is None + + +def test_full_text_reader_keeps_query_and_limit_captured_before_transport(search_table): + builder = _builder(search_table, 'text') + native, _, _ = _native_builder('text') + with patch(_native_factory('text'), return_value=native): + reader = builder.new_full_text_read() + builder.with_query('body', 'changed').with_limit(99) + transported = pickle.loads(pickle.dumps(reader)) + with _forbid_python_execution(), patch(_native_factory('text'), return_value=native) as prepare, patch( + 'pypaimon.execution.native.read.native_search_plan._plan_to_native', return_value=object()): + transported.read_plan(_ordinary_plan('text')) + captured = prepare.call_args[0][0] + assert (captured._field_name, captured._query, captured._limit) == ('body', 'needle', 2) + assert transported._python_executor is None @pytest.mark.parametrize('name', [ @@ -221,7 +272,7 @@ def test_vector_split_uses_its_shared_contract_module(): assert pickle.loads(pickle.dumps(VectorSearchSplit)) is VectorSearchSplit -def test_single_reader_freezes_query_values_before_transport_fallback(search_table): +def test_single_reader_freezes_query_values_before_transport(search_table): query = [1.0, 0.0] builder = _builder(search_table, 'vector').with_query_vector(query) native, _, _ = _native_builder('vector') @@ -229,5 +280,110 @@ def test_single_reader_freezes_query_values_before_transport_fallback(search_tab reader = builder.new_vector_search_read() query[:] = [3.0, 4.0] assert native.with_query_vector.call_args[0][0] == [1.0, 0.0] - fallback = search_dispatch.python_search_reader(pickle.loads(pickle.dumps(reader))) - assert fallback._query_vector == [1.0, 0.0] + transported = pickle.loads(pickle.dumps(reader)) + with _forbid_python_execution(), patch(_native_factory('vector'), return_value=native), patch( + 'pypaimon.execution.native.read.native_search_plan._plan_to_native', return_value=object()): + transported.read_plan(_ordinary_plan('vector')) + assert native.with_query_vector.call_args[0][0] == [1.0, 0.0] + assert transported._python_executor is None + + +@pytest.mark.parametrize('family', ['vector', 'batch', 'text']) +def test_python_source_metadata_round_trips_through_public_native_plan_constructor(search_table, family): + from pypaimon.execution.native.read.native_search_plan import _plan_from_native, _plan_to_native + from pypaimon.index.index_file_meta import IndexFileMeta + from pypaimon.snapshot.snapshot import Snapshot + from pypaimon.table.source.full_text_search_split import IndexFullTextSearchSplit, RawFullTextSearchSplit + from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit, RawVectorSearchSplit + from pypaimon.utils.range import Range + + class Plan: + def __init__(self, snapshot_json, splits_json): + self.snapshot_json = snapshot_json + self.splits_json = splits_json + + def metadata(self): + snapshot = None if self.snapshot_json is None else json.loads(self.snapshot_json) + return json.dumps({'snapshot': snapshot, 'snapshot_id': None if snapshot is None else snapshot['id'], + 'splits': json.loads(self.splits_json)}) + + sdk, datafusion = ModuleType('pypaimon_rust'), ModuleType('pypaimon_rust.datafusion') + datafusion.VectorScanPlan = Plan + datafusion.FullTextScanPlan = Plan + sdk.datafusion = datafusion + full_text = family == 'text' + snapshot = Snapshot(3, 7, 2, 'base', 'delta', 10, 10, 'writer', 1, 'APPEND', 0, next_row_id=10) + index = IndexFileMeta('full-text' if full_text else 'ivf-flat', 'index-1', 1024, 4, + external_path='s3://bucket/index-1') + scalar = IndexFileMeta('btree', 'index-scalar', 128, 4) + splits = ([IndexFullTextSearchSplit('body', 0, 3, [index], [scalar]), RawFullTextSearchSplit([Range(4, 9)])] + if full_text else [IndexVectorSearchSplit(0, 3, [index], [scalar]), + RawVectorSearchSplit([Range(4, 9)], [scalar], 'ivf-flat')]) + source = (FullTextScanPlan if full_text else VectorSearchScanPlan)(splits, snapshot) + with patch.dict('sys.modules', {'pypaimon_rust': sdk, 'pypaimon_rust.datafusion': datafusion}): + native = _plan_to_native(source, full_text, 'embedding') + restored = _plan_from_native(search_table, native, full_text, False) + copied = _plan_to_native(pickle.loads(pickle.dumps(restored)), full_text, 'embedding') + empty = _plan_to_native(_ordinary_plan(family), full_text, 'embedding') + assert json.loads(native.snapshot_json)['schemaId'] == 2 + assert json.loads(native.splits_json)[0]['column'] == ('body' if full_text else 'embedding') + assert copied.metadata() == native.metadata() + assert restored.snapshot() == snapshot + restored_index = restored.splits()[0] + restored_files = restored_index.full_text_index_files if full_text else restored_index.vector_index_files + assert restored_files[0].external_path == index.external_path + assert restored_index.scalar_index_files == [scalar] + assert json.loads(empty.metadata()) == {'snapshot': None, 'snapshot_id': None, 'splits': []} + assert not hasattr(restored, '_native_plan') + + +@pytest.mark.parametrize('full_text', [False, True]) +def test_primary_key_plan_with_only_snapshot_id_resolves_that_exact_snapshot(full_text): + from pypaimon.execution.native.read.native_search_plan import NativeSearchRead + from pypaimon.table.source.search_types import PrimaryKeyFullTextScanPlan, PrimaryKeyVectorScanPlan + + snapshot = SimpleNamespace(id=7) + manager = SimpleNamespace( + get_snapshot_by_id=Mock(return_value=snapshot), + get_latest_snapshot=Mock(side_effect=AssertionError('Resolved latest instead of the requested snapshot'))) + table = SimpleNamespace(snapshot_manager=Mock(return_value=manager)) + core_result = SimpleNamespace(snapshot_id=lambda: 7, positions=lambda: [], splits=lambda: []) + core = SimpleNamespace(read_plan=Mock(return_value=core_result)) + reader = NativeSearchRead(table, core, True, full_text=full_text) + sources = [object()] + requested = (PrimaryKeyFullTextScanPlan if full_text else PrimaryKeyVectorScanPlan)(7, sources) + native_plan = object() + with patch('pypaimon.execution.native.read.native_search_plan._plan_to_native', + return_value=native_plan) as convert, patch( + 'pypaimon.snapshot.time_travel_util.TimeTravelUtil.resolve_snapshot', + side_effect=AssertionError('Resolved table time travel instead of the explicit plan snapshot')): + result = reader.read_plan(requested) + manager.get_snapshot_by_id.assert_called_once_with(7) + manager.get_latest_snapshot.assert_not_called() + converted = convert.call_args[0][0] + assert converted.snapshot() is snapshot + assert converted.splits() is sources + assert requested.snapshot() is None + assert result.snapshot_id == 7 + core.read_plan.assert_called_once_with(native_plan) + + +@pytest.mark.parametrize('full_text', [False, True]) +def test_missing_primary_key_plan_snapshot_does_not_fall_back_to_latest(full_text): + from pypaimon.execution.native.read.native_search_plan import NativeSearchRead + from pypaimon.table.source.search_types import PrimaryKeyFullTextScanPlan, PrimaryKeyVectorScanPlan + + manager = SimpleNamespace( + get_snapshot_by_id=Mock(return_value=None), + get_latest_snapshot=Mock(side_effect=AssertionError('Expired plans cannot be rebound to latest'))) + table = SimpleNamespace(snapshot_manager=Mock(return_value=manager)) + core = SimpleNamespace(read_plan=Mock()) + reader = NativeSearchRead(table, core, True, full_text=full_text) + requested = (PrimaryKeyFullTextScanPlan if full_text else PrimaryKeyVectorScanPlan)(7, [object()]) + with patch('pypaimon.execution.native.read.native_search_plan._plan_to_native') as convert: + with pytest.raises(ValueError, match='7'): + reader.read_plan(requested) + manager.get_snapshot_by_id.assert_called_once_with(7) + manager.get_latest_snapshot.assert_not_called() + core.read_plan.assert_not_called() + convert.assert_not_called() diff --git a/paimon-python/pypaimon/tests/native_search_plan_test.py b/paimon-python/pypaimon/tests/native_search_plan_test.py index 0d43e591d3c3..3c84d473d468 100644 --- a/paimon-python/pypaimon/tests/native_search_plan_test.py +++ b/paimon-python/pypaimon/tests/native_search_plan_test.py @@ -138,26 +138,22 @@ def test_explicit_native_plan_and_read_avoid_python_search(table, family, mode, @pytest.mark.parametrize('family', ['vector', 'batch', 'full-text']) -@pytest.mark.parametrize('transport', ['pickle-plan', 'pickle-reader', 'split-list']) +@pytest.mark.parametrize('transport', ['pickle-plan', 'pickle-reader', 'pickle-both', 'python-plan']) def test_transported_plan_uses_captured_snapshot_without_replanning(table, family, transport): native = _builder(table, family, native=True, limit=20) - plan = _scan(native, family).scan() + planner = _builder(table, family, limit=20) if transport == 'python-plan' else native + plan = _scan(planner, family).scan() reader = _reader(native, family) expected = _read(reader, family, plan) snapshot = plan.snapshot().id captured = pickle.dumps(plan.splits()) indexes._append(table, hybrid._rows(12, 4), hybrid.SCHEMA) - if transport == 'pickle-plan': + if transport in ('pickle-plan', 'pickle-both'): plan = pickle.loads(pickle.dumps(plan)) - elif transport == 'pickle-reader': + if transport in ('pickle-reader', 'pickle-both'): reader = pickle.loads(pickle.dumps(reader)) - with patch.object(DataEvolutionVectorScan, 'scan', side_effect=AssertionError('Replanned vector search')), \ - patch.object(DataEvolutionFullTextScan, 'scan', side_effect=AssertionError('Replanned full-text search')): - if transport == 'split-list': - classic = _reader(_builder(table._copy_with_snapshot(plan.snapshot()), family, limit=20), family) - actual = classic.read_batch(plan.splits()) if family == 'batch' else [classic.read(plan.splits())] - else: - actual = _read(reader, family, plan) + with _native_only(): + actual = _read(reader, family, plan) assert plan.snapshot().id == snapshot assert pickle.dumps(plan.splits()) == captured _assert_results(actual, expected) @@ -169,17 +165,37 @@ def test_modified_plan_and_shallow_alias_observe_current_split_selection(table, plan = _scan(builder, family).scan() alias = copy(plan) alias.splits().clear() - results = _read(_reader(builder, family), family, plan) + with _native_only(): + results = _read(_reader(builder, family), family, plan) assert len(results) == (2 if family == 'batch' else 1) assert all(result.is_empty() for result in results) @pytest.mark.parametrize('family', ['vector', 'batch', 'full-text']) -def test_read_plan_context_errors_do_not_retry_python(table, family): - plan = _scan(_builder(table, family, native=True), family).scan() - reader = _reader(_builder(table, family, native=True, filtered=True), family) - with _native_only(), pytest.raises(Exception, match='same table|same.*filter|different table|different.*filter'): - _read(reader, family, plan) +def test_subset_of_python_sources_remains_authoritative_for_native(table, family): + classic = _builder(table, family, limit=20) + plan = _scan(classic, family).scan() + assert len(plan.splits()) > 1 + plan.splits()[:] = plan.splits()[:1] + expected = _read(_reader(classic, family), family, plan) + sources = pickle.dumps(plan.splits()) + reader = _reader(_builder(table, family, native=True, limit=20), family) + with _native_only(): + actual = _read(pickle.loads(pickle.dumps(reader)), family, pickle.loads(pickle.dumps(plan))) + _assert_results(actual, expected) + assert pickle.dumps(plan.splits()) == sources + + +@pytest.mark.parametrize('family', ['vector', 'batch', 'full-text']) +def test_reader_filter_applies_to_explicit_sources_without_replanning(table, family): + plan = _scan(_builder(table, family, native=True, limit=20), family).scan() + captured_sources = pickle.dumps(plan.splits()) + expected = _read(_reader(_builder(table, family, filtered=True, limit=20), family), family, plan) + reader = _reader(_builder(table, family, native=True, filtered=True, limit=20), family) + with _native_only(): + actual = _read(reader, family, pickle.loads(pickle.dumps(plan))) + _assert_results(actual, expected) + assert pickle.dumps(plan.splits()) == captured_sources @pytest.mark.parametrize('family', ['vector', 'batch', 'full-text']) @@ -282,7 +298,9 @@ def test_empty_native_plans_remain_transportable(rest_catalog, family): assert plan.snapshot() is None assert plan.splits() == [] assert all(result.is_empty() for result in results) - _assert_results(_read(_reader(builder, family), family, pickle.loads(pickle.dumps(plan))), results) + with _native_only(): + copied_reader = pickle.loads(pickle.dumps(_reader(builder, family))) + _assert_results(_read(copied_reader, family, pickle.loads(pickle.dumps(plan))), results) def test_vector_scan_does_not_require_query_or_limit(table): @@ -306,11 +324,13 @@ def test_full_text_scan_does_not_require_limit(table): @pytest.mark.parametrize('family', ['vector', 'batch', 'full-text']) -def test_native_plan_is_query_independent_and_pickle_keeps_original_handle(table, family): +def test_native_plan_is_query_independent_and_pickle_keeps_plain_sources(table, family): builder = _builder(table, family, native=True, limit=1) plan = _scan(builder, family).scan() - pickle.dumps(plan) - assert plan._unchanged_native_plan() is not None + captured = pickle.dumps(plan) + assert not hasattr(plan, '_native_plan') + assert not hasattr(plan, '_native_fingerprint') + assert pickle.dumps(pickle.loads(captured)) == captured if family == 'full-text': builder.with_query('content', text.match_query('lake')) elif family == 'batch': @@ -410,7 +430,9 @@ def test_pk_plan_retains_snapshot_and_sources_after_primary_key_update(source, p assert table.snapshot_manager().get_latest_snapshot().id > captured with _native_only(): actual = _read(_reader(builder, family), family, plan) - transported = _read(_reader(builder, family), family, pickle.loads(pickle.dumps(plan))) + with _native_only(): + transported = _read(pickle.loads(pickle.dumps(_reader(builder, family))), family, + pickle.loads(pickle.dumps(plan))) _assert_results(actual, expected) _assert_results(transported, expected) assert vectors._rows(table, actual[0]).to_pylist() == expected_rows diff --git a/paimon-python/pypaimon/tests/search_task_execution_test.py b/paimon-python/pypaimon/tests/search_task_execution_test.py index bc90e00de2d9..2029169fa665 100644 --- a/paimon-python/pypaimon/tests/search_task_execution_test.py +++ b/paimon-python/pypaimon/tests/search_task_execution_test.py @@ -32,6 +32,15 @@ from pypaimon.table.source.hybrid_search_builder import HybridSearchBuilderImpl +@pytest.fixture(autouse=True) +def _native_plan_conversion(): + # These tests exercise callback ownership and cancellation. Source codec + # construction has dedicated tests and must not require an installed SDK. + with patch('pypaimon.execution.native.read.native_search_plan._plan_to_native', + side_effect=lambda plan, full_text=False, column=None: plan): + yield + + class InlineExecutor(SearchTaskExecutor): def __init__(self): self.tasks = [] @@ -76,9 +85,8 @@ def read(native_plan, executor): assert list(executor(b'shared-context', [b'index-task', b'raw-task'])) == [b'index-result', b'raw-result'] return [result, result] if family == 'batch' else result - core_read = SimpleNamespace(read_plan=read, read_batch_plan=read) - native = NativeSearchRead(table, core_read, False, ()) - native.resolve_plan = lambda requested: plan + core_read = SimpleNamespace(_read_plan_with_executor=read, _read_batch_plan_with_executor=read) + native = NativeSearchRead(table, core_read, False, full_text=family == 'full-text', batch=family == 'batch') python_factory = Mock(side_effect=AssertionError('Python executor was constructed')) reader = (_BatchSearchRead if family == 'batch' else _SplitSearchRead)(native, python_factory) @@ -127,7 +135,7 @@ def close(self): def test_distributed_native_failure_does_not_start_python_search(): failure = CancelledError('remote task failed') - native = SimpleNamespace(resolve_plan=lambda plan: plan, read_plan=Mock(side_effect=failure)) + native = SimpleNamespace(read_plan=Mock(side_effect=failure)) fallback = Mock(side_effect=AssertionError('Python search replayed')) reader = _SplitSearchRead(native, fallback) with pytest.raises(CancelledError) as caught: @@ -146,7 +154,7 @@ def read(plan, executor): callbacks.append(executor) raise failure - reader = NativeSearchRead(table, SimpleNamespace(read_plan=read), False, ()) + reader = NativeSearchRead(table, SimpleNamespace(_read_plan_with_executor=read), False) executor = InlineExecutor() with pytest.raises(ValueError) as caught: reader.read_plan(object(), executor=executor) @@ -190,7 +198,7 @@ def read(plan, executor): table = SimpleNamespace(options=SimpleNamespace( options=SimpleNamespace(to_map=lambda: {}), data_evolution_enabled=lambda: True)) - reader = NativeSearchRead(table, SimpleNamespace(read_plan=read), False, ()) + reader = NativeSearchRead(table, SimpleNamespace(_read_plan_with_executor=read), False) with pytest.raises(ValueError) as caught: reader.read_plan(object(), executor=MapExecutor(lambda tasks, ordered: FailedResults())) assert caught.value is failure From e33b19df441878ce543d170f2db0b9445a3b92a3 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Wed, 21 Oct 2026 22:47:08 +0800 Subject: [PATCH 3/9] Adapt native search plans through typed snapshot and split objects --- docs/docs/pypaimon/multimodal-search.md | 16 +- paimon-python/pypaimon/execution/README.md | 18 +- .../native/read/native_search_plan.py | 113 ++++---- .../tests/lazy_search_execution_test.py | 273 ++++++++++++++++-- .../pypaimon/tests/native_search_plan_test.py | 25 ++ .../tests/search_task_execution_test.py | 2 +- 6 files changed, 355 insertions(+), 92 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-search.md b/docs/docs/pypaimon/multimodal-search.md index 30544083259e..0993d90521bb 100644 --- a/docs/docs/pypaimon/multimodal-search.md +++ b/docs/docs/pypaimon/multimodal-search.md @@ -359,9 +359,11 @@ results = builder.new_batch_vector_search_read().read_batch_plan( plan, executor=RaySearchTaskExecutor(concurrency=4)) ``` -Plans contain ordinary snapshot and split metadata. A copied, pickled, or -Python-constructed plan can be passed to a Native reader without choosing the -Python backend. For example, to execute an explicitly selected set of sources: +Plans contain ordinary snapshot and typed split metadata. A copied, pickled, +or Python-constructed plan can be passed to a Native reader without choosing +the Python backend. The Native SDK also accepts snapshot and split objects +through its plan constructors. For example, to execute an explicitly selected +set of sources: ```python from pypaimon.table.source.search_types import VectorSearchScanPlan @@ -371,9 +373,11 @@ results = builder.new_batch_vector_search_read().read_batch_plan(selected) ``` The supplied splits are authoritative: readers do not add uncovered raw ranges -or rediscover other search indexes. Query vectors, limits and filters belong to -the reader. Rust Core coordinates global candidate selection and refinement; -its internal shard protocol is separate from the public source plan. +or rediscover other search indexes. A plan with no captured snapshot stays +empty rather than selecting the latest snapshot. Query vectors, limits and +filters belong to the reader. Rust Core coordinates global candidate selection +and refinement; its internal shard protocol is separate from the public source +plan. Hybrid builders expose `execute(executor=...)` for the same scheduler. The existing `execute_local()` methods retain their local defaults. diff --git a/paimon-python/pypaimon/execution/README.md b/paimon-python/pypaimon/execution/README.md index 5ae09dd3136d..8e9a33138097 100644 --- a/paimon-python/pypaimon/execution/README.md +++ b/paimon-python/pypaimon/execution/README.md @@ -154,13 +154,17 @@ opaque shared context and task bytes; Core owns both their interpretation and global query coordination. The global candidate limit is applied before refinement, and full-text raw fallback retains one corpus for BM25 statistics. -Vector and full-text plans are snapshot-and-split data in both backends. Rust -plans have public constructors and can be serialized; the Native adapter builds -them from ordinary Python plans, including plans made by Python scans or sent -to another process. Readers own the table, schema, query and filter. Native -selection is retained after transport, and source subsets never add missing -indexes or raw ranges. Shard callbacks use an internal binding bridge; they are -not part of Rust's public Plan/Read API. +Vector and full-text plans are snapshot-and-split data in both backends. The +Rust SDK constructors accept typed `Snapshot` and search split objects; indexed, +raw and primary-key splits keep their own fields and index metadata. The Native +adapter translates ordinary Python plans into these values, including plans +made by Python scans or sent to another process. Existing snapshot and index +metadata codecs are private adapters, rather than a public JSON plan protocol. +Readers own the table, schema, query and filter. Native selection is retained +after transport, and source subsets never add missing indexes or raw ranges. +An absent snapshot describes an empty capture and never selects latest. Shard +callbacks use an internal binding bridge; they are not part of Rust's public +Plan/Read API. ## Corresponding execution APIs diff --git a/paimon-python/pypaimon/execution/native/read/native_search_plan.py b/paimon-python/pypaimon/execution/native/read/native_search_plan.py index a470957f696c..95cebb37878c 100644 --- a/paimon-python/pypaimon/execution/native/read/native_search_plan.py +++ b/paimon-python/pypaimon/execution/native/read/native_search_plan.py @@ -17,7 +17,6 @@ """Translate ordinary snapshot-and-split plans to and from Rust Core.""" -import json from copy import copy, deepcopy from pypaimon.common.json_util import JSON @@ -36,47 +35,53 @@ def _range(value): - return Range(value['from'], value['to']) + return Range(*value) def _files(values): - return [_index_from_row(bytes(value)) for value in values] + return [_index_from_row(bytes(value._serialize())) for value in values] -def _plan_from_native(table, native, full_text, primary_key): +def _split_from_native(table, split, full_text): + from pypaimon_rust import datafusion as sdk from pypaimon.table.source.full_text_search_split import IndexFullTextSearchSplit, RawFullTextSearchSplit from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit, RawVectorSearchSplit - metadata = json.loads(native.metadata()) - snapshot = JSON.from_json(json.dumps(metadata['snapshot']), Snapshot) if metadata['snapshot'] is not None else None - splits = [] - for value in metadata['splits']: - if value['kind'] == 'primary-key': - split = deserialize_split_v1(bytes(value['data_split']), _partition_fields(table), - table.trimmed_primary_keys_fields) - ranges = {name: tuple(_range(range_) for range_ in ranges) - for name, ranges in value['row_ranges_by_file'].items()} - cls = PrimaryKeyFullTextSearchSplit if full_text else PrimaryKeyVectorSearchSplit - splits.append(cls(split, tuple(_files(value['payloads'])), tuple(value['uncovered_data_files']), ranges)) - elif value['kind'] == 'index': - row_range = _range(value['row_range']) - args = (row_range.from_, row_range.to, _files(value['index_files']), _files(value['scalar_index_files'])) - splits.append(IndexFullTextSearchSplit(value['column'], *args) - if full_text else IndexVectorSearchSplit(*args)) - else: - ranges = [_range(range_) for range_ in value['row_ranges']] - splits.append(RawFullTextSearchSplit(ranges) if full_text else RawVectorSearchSplit( - ranges, _files(value['scalar_index_files']), value['index_type'])) + native_pk = sdk.PrimaryKeyFullTextSearchSplit if full_text else sdk.PrimaryKeyVectorSearchSplit + native_index = sdk.IndexFullTextSearchSplit if full_text else sdk.IndexVectorSearchSplit + native_raw = sdk.RawFullTextSearchSplit if full_text else sdk.RawVectorSearchSplit + if isinstance(split, native_pk): + data_split = deserialize_split_v1(bytes(split.data_split().serialize()), _partition_fields(table), + table.trimmed_primary_keys_fields) + ranges = {name: tuple(_range(value) for value in values) + for name, values in split.row_ranges_by_file().items()} + cls = PrimaryKeyFullTextSearchSplit if full_text else PrimaryKeyVectorSearchSplit + return cls(data_split, tuple(_files(split.payloads())), tuple(split.uncovered_data_files()), ranges) + if isinstance(split, native_index): + args = (split.row_range_start(), split.row_range_end(), + _files(split.full_text_index_files() if full_text else split.vector_index_files()), + _files(split.scalar_index_files())) + return IndexFullTextSearchSplit(split.column_name(), *args) if full_text else IndexVectorSearchSplit(*args) + if isinstance(split, native_raw): + ranges = [_range(value) for value in split.row_ranges()] + return RawFullTextSearchSplit(ranges) if full_text else RawVectorSearchSplit( + ranges, _files(split.scalar_index_files()), split.index_type()) + raise TypeError('Unsupported Native %s search split: %s' % ( + 'full-text' if full_text else 'vector', type(split).__name__)) + + +def _plan_from_native(table, native, full_text, primary_key): + captured = native.snapshot() + snapshot = None if captured is None else JSON.from_json(captured._to_json(), Snapshot) + splits = [_split_from_native(table, split, full_text) for split in native.splits()] if primary_key: cls = PrimaryKeyFullTextScanPlan if full_text else PrimaryKeyVectorScanPlan - plan = cls(metadata['snapshot_id'] or 0, splits, snapshot) - else: - cls = FullTextScanPlan if full_text else VectorSearchScanPlan - plan = cls(splits, snapshot) - return plan + return cls(0 if snapshot is None else snapshot.id, splits, snapshot) + return (FullTextScanPlan if full_text else VectorSearchScanPlan)(splits, snapshot) -def _split_metadata(split, full_text, column): +def _split_to_native(split, full_text): + from pypaimon_rust import datafusion as sdk from pypaimon.table.source.full_text_search_split import IndexFullTextSearchSplit, RawFullTextSearchSplit from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit, RawVectorSearchSplit @@ -87,37 +92,34 @@ def _split_metadata(split, full_text, column): 'full-text' if full_text else 'vector', type(split).__name__)) def files(values): - return [list(_index_row(value)) for value in values] + return [sdk.IndexFileMeta._deserialize(_index_row(value)) for value in values] - def range_(value): - return {'from': value.from_, 'to': value.to} + def ranges(values): + return [(value.from_, value.to) for value in values] if isinstance(split, (PrimaryKeyVectorSearchSplit, PrimaryKeyFullTextSearchSplit)): - return {'kind': 'primary-key', 'data_split': list(serialize_split_v1(split.data_split)), - 'payloads': files(split.payloads), 'uncovered_data_files': list(split.uncovered_data_files), - 'row_ranges_by_file': {name: [range_(value) for value in values] - for name, values in split.row_ranges_by_file.items()}} - if isinstance(split, (IndexVectorSearchSplit, IndexFullTextSearchSplit)): - return {'kind': 'index', 'column': split.column_name if full_text else column, - 'row_range': {'from': split.row_range_start, 'to': split.row_range_end}, - 'index_files': files(split.full_text_index_files if full_text else split.vector_index_files), - 'scalar_index_files': files(split.scalar_index_files)} + cls = sdk.PrimaryKeyFullTextSearchSplit if full_text else sdk.PrimaryKeyVectorSearchSplit + return cls(sdk.Split.deserialize(serialize_split_v1(split.data_split)), files(split.payloads), + list(split.uncovered_data_files), + {name: ranges(values) for name, values in split.row_ranges_by_file.items()}) + if isinstance(split, IndexFullTextSearchSplit): + return sdk.IndexFullTextSearchSplit(split.column_name, split.row_range_start, split.row_range_end, + files(split.full_text_index_files), files(split.scalar_index_files)) + if isinstance(split, IndexVectorSearchSplit): + return sdk.IndexVectorSearchSplit(split.row_range_start, split.row_range_end, + files(split.vector_index_files), files(split.scalar_index_files)) if isinstance(split, RawVectorSearchSplit): - return {'kind': 'raw', 'row_ranges': [range_(value) for value in split.row_ranges], - 'scalar_index_files': files(split.scalar_index_files), 'index_type': split.index_type} - if isinstance(split, RawFullTextSearchSplit): - return {'kind': 'raw', 'row_ranges': [range_(value) for value in split.row_ranges], - 'scalar_index_files': [], 'index_type': None} - raise TypeError('Unsupported search split: %s' % type(split).__name__) + return sdk.RawVectorSearchSplit(ranges(split.row_ranges), files(split.scalar_index_files), split.index_type) + return sdk.RawFullTextSearchSplit(ranges(split.row_ranges)) -def _plan_to_native(plan, full_text, column=None): - from pypaimon_rust.datafusion import FullTextScanPlan as NativeFullTextPlan, VectorScanPlan as NativeVectorPlan +def _plan_to_native(plan, full_text): + from pypaimon_rust import datafusion as sdk - snapshot = plan.snapshot() - return (NativeFullTextPlan if full_text else NativeVectorPlan)( - None if snapshot is None else JSON.to_json(snapshot), - json.dumps([_split_metadata(split, full_text, column) for split in plan.splits()])) + captured = plan.snapshot() + snapshot = None if captured is None else sdk.Snapshot._from_json(JSON.to_json(captured)) + cls = sdk.FullTextScanPlan if full_text else sdk.VectorScanPlan + return cls([_split_to_native(split, full_text) for split in plan.splits()], snapshot) def _capture_builder(builder): @@ -192,8 +194,7 @@ def _plan(self, plan): raise ValueError('Captured search snapshot %s is unavailable' % plan.snapshot_id) plan = copy(plan) plan._snapshot = snapshot - column = None if self._builder is None or self._full_text else self._builder._vector_column.name - return _plan_to_native(plan, self._full_text, column) + return _plan_to_native(plan, self._full_text) def read_plan(self, plan, executor=None): native = self._get_native_reader() diff --git a/paimon-python/pypaimon/tests/lazy_search_execution_test.py b/paimon-python/pypaimon/tests/lazy_search_execution_test.py index 1eeb3830f480..4e0b5b6799ff 100644 --- a/paimon-python/pypaimon/tests/lazy_search_execution_test.py +++ b/paimon-python/pypaimon/tests/lazy_search_execution_test.py @@ -60,7 +60,7 @@ def _factories(builder, family): def _native_builder(family): - plan = SimpleNamespace(metadata=lambda: '{"snapshot":null,"snapshot_id":null,"splits":[]}') + plan = SimpleNamespace(snapshot=lambda: None, splits=lambda: []) result = SimpleNamespace(row_ids=lambda: {7: 0.5}) reader = SimpleNamespace(read_plan=Mock(return_value=result), read_batch_plan=Mock(return_value=[result, result])) scan = SimpleNamespace(scan=Mock(return_value=plan)) @@ -288,6 +288,150 @@ def test_single_reader_freezes_query_values_before_transport(search_table): assert transported._python_executor is None +def _typed_native_sdk(): + """Strict SDK shapes exercise object transport without loading an extension.""" + class NativeSnapshot: + @classmethod + def _from_json(cls, value): + assert isinstance(value, str) + result = cls() + result.value = json.loads(value) + return result + + def _to_json(self): + return json.dumps(self.value) + + def id(self): + return self.value['id'] + + class NativeIndexFileMeta: + @classmethod + def _deserialize(cls, value): + assert isinstance(value, bytes) + result = cls() + result.value = value + return result + + def _serialize(self): + return self.value + + class NativeSplit: + @classmethod + def deserialize(cls, value): + assert isinstance(value, bytes) + result = cls() + result.value = value + return result + + def serialize(self): + return self.value + + class IndexSplit: + def __init__(self, start, end, files, scalar=()): + assert all(isinstance(file, NativeIndexFileMeta) for file in list(files) + list(scalar)) + self.values = start, end, list(files), list(scalar) + + def row_range_start(self): + return self.values[0] + + def row_range_end(self): + return self.values[1] + + def scalar_index_files(self): + return self.values[3] + + class IndexVectorSplit(IndexSplit): + def vector_index_files(self): + return self.values[2] + + class IndexTextSplit(IndexSplit): + def __init__(self, column, start, end, files, scalar=()): + super().__init__(start, end, files, scalar) + self.column = column + + def column_name(self): + return self.column + + def full_text_index_files(self): + return self.values[2] + + class RawVectorSplit: + def __init__(self, ranges, scalar=(), index_type=None): + assert all(isinstance(value, tuple) and len(value) == 2 for value in ranges) + self.values = list(ranges), list(scalar), index_type + + def row_ranges(self): + return self.values[0] + + def scalar_index_files(self): + return self.values[1] + + def index_type(self): + return self.values[2] + + class RawTextSplit: + def __init__(self, ranges): + assert all(isinstance(value, tuple) and len(value) == 2 for value in ranges) + self.ranges = list(ranges) + + def row_ranges(self): + return self.ranges + + class PrimaryKeySplit: + def __init__(self, split, payloads, uncovered, ranges): + assert isinstance(split, NativeSplit) + assert all(isinstance(file, NativeIndexFileMeta) for file in payloads) + self.values = split, list(payloads), list(uncovered), dict(ranges) + + def data_split(self): + return self.values[0] + + def payloads(self): + return self.values[1] + + def uncovered_data_files(self): + return self.values[2] + + def row_ranges_by_file(self): + return self.values[3] + + class PrimaryKeyVectorSplit(PrimaryKeySplit): + pass + + class PrimaryKeyTextSplit(PrimaryKeySplit): + pass + + class Plan: + def __init__(self, splits, snapshot=None): + assert isinstance(splits, (list, tuple)) + assert snapshot is None or isinstance(snapshot, NativeSnapshot) + assert all(type(split) in self.split_types for split in splits) + self.sources, self.captured = list(splits), snapshot + + def splits(self): + return self.sources + + def snapshot(self): + return self.captured + + class VectorPlan(Plan): + split_types = IndexVectorSplit, RawVectorSplit, PrimaryKeyVectorSplit + + class TextPlan(Plan): + split_types = IndexTextSplit, RawTextSplit, PrimaryKeyTextSplit + + sdk, datafusion = ModuleType('pypaimon_rust'), ModuleType('pypaimon_rust.datafusion') + for name, cls in [('Snapshot', NativeSnapshot), ('IndexFileMeta', NativeIndexFileMeta), ('Split', NativeSplit), + ('VectorScanPlan', VectorPlan), ('FullTextScanPlan', TextPlan), + ('IndexVectorSearchSplit', IndexVectorSplit), ('RawVectorSearchSplit', RawVectorSplit), + ('PrimaryKeyVectorSearchSplit', PrimaryKeyVectorSplit), + ('IndexFullTextSearchSplit', IndexTextSplit), ('RawFullTextSearchSplit', RawTextSplit), + ('PrimaryKeyFullTextSearchSplit', PrimaryKeyTextSplit)]: + setattr(datafusion, name, cls) + sdk.datafusion = datafusion + return sdk, datafusion + + @pytest.mark.parametrize('family', ['vector', 'batch', 'text']) def test_python_source_metadata_round_trips_through_public_native_plan_constructor(search_table, family): from pypaimon.execution.native.read.native_search_plan import _plan_from_native, _plan_to_native @@ -297,20 +441,7 @@ def test_python_source_metadata_round_trips_through_public_native_plan_construct from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit, RawVectorSearchSplit from pypaimon.utils.range import Range - class Plan: - def __init__(self, snapshot_json, splits_json): - self.snapshot_json = snapshot_json - self.splits_json = splits_json - - def metadata(self): - snapshot = None if self.snapshot_json is None else json.loads(self.snapshot_json) - return json.dumps({'snapshot': snapshot, 'snapshot_id': None if snapshot is None else snapshot['id'], - 'splits': json.loads(self.splits_json)}) - - sdk, datafusion = ModuleType('pypaimon_rust'), ModuleType('pypaimon_rust.datafusion') - datafusion.VectorScanPlan = Plan - datafusion.FullTextScanPlan = Plan - sdk.datafusion = datafusion + sdk, datafusion = _typed_native_sdk() full_text = family == 'text' snapshot = Snapshot(3, 7, 2, 'base', 'delta', 10, 10, 'writer', 1, 'APPEND', 0, next_row_id=10) index = IndexFileMeta('full-text' if full_text else 'ivf-flat', 'index-1', 1024, 4, @@ -321,19 +452,24 @@ def metadata(self): RawVectorSearchSplit([Range(4, 9)], [scalar], 'ivf-flat')]) source = (FullTextScanPlan if full_text else VectorSearchScanPlan)(splits, snapshot) with patch.dict('sys.modules', {'pypaimon_rust': sdk, 'pypaimon_rust.datafusion': datafusion}): - native = _plan_to_native(source, full_text, 'embedding') + native = _plan_to_native(source, full_text) restored = _plan_from_native(search_table, native, full_text, False) - copied = _plan_to_native(pickle.loads(pickle.dumps(restored)), full_text, 'embedding') - empty = _plan_to_native(_ordinary_plan(family), full_text, 'embedding') - assert json.loads(native.snapshot_json)['schemaId'] == 2 - assert json.loads(native.splits_json)[0]['column'] == ('body' if full_text else 'embedding') - assert copied.metadata() == native.metadata() + copied = _plan_to_native(pickle.loads(pickle.dumps(restored)), full_text) + empty = _plan_to_native(_ordinary_plan(family), full_text) + assert json.loads(native.snapshot()._to_json())['schemaId'] == 2 + assert isinstance(native.snapshot(), datafusion.Snapshot) + assert isinstance(native.splits()[0], datafusion.IndexFullTextSearchSplit if full_text else + datafusion.IndexVectorSearchSplit) + assert copied.snapshot()._to_json() == native.snapshot()._to_json() + assert type(copied.splits()[0]) is type(native.splits()[0]) assert restored.snapshot() == snapshot restored_index = restored.splits()[0] restored_files = restored_index.full_text_index_files if full_text else restored_index.vector_index_files assert restored_files[0].external_path == index.external_path assert restored_index.scalar_index_files == [scalar] - assert json.loads(empty.metadata()) == {'snapshot': None, 'snapshot_id': None, 'splits': []} + assert restored.splits() == source.splits() + assert empty.snapshot() is None and empty.splits() == [] + assert not hasattr(native, 'metadata') and not hasattr(native, 'snapshot_id') assert not hasattr(restored, '_native_plan') @@ -387,3 +523,96 @@ def test_missing_primary_key_plan_snapshot_does_not_fall_back_to_latest(full_tex manager.get_latest_snapshot.assert_not_called() core.read_plan.assert_not_called() convert.assert_not_called() + + +@pytest.mark.parametrize('full_text', [False, True]) +def test_primary_key_sources_round_trip_as_typed_splits(search_table, full_text): + from pypaimon.execution.native.read.native_search_plan import _plan_from_native, _plan_to_native + from pypaimon.index.index_file_meta import IndexFileMeta + from pypaimon.manifest.schema.data_file_meta import DataFileMeta + from pypaimon.manifest.schema.simple_stats import SimpleStats + from pypaimon.read.split import DataSplit + from pypaimon.read.split_serializer import serialize_split_v1 + from pypaimon.snapshot.snapshot import Snapshot + from pypaimon.table.row.generic_row import GenericRow + from pypaimon.table.source.deletion_file import DeletionFile + from pypaimon.table.source.search_types import ( + PrimaryKeyFullTextScanPlan, PrimaryKeyFullTextSearchSplit, + PrimaryKeyVectorScanPlan, PrimaryKeyVectorSearchSplit) + from pypaimon.utils.range import Range + + sdk, datafusion = _typed_native_sdk() + snapshot = Snapshot(3, 7, 2, 'base', 'delta', 10, 10, 'writer', 1, 'APPEND', 0, next_row_id=10) + empty = GenericRow([], []) + file = DataFileMeta.create('tail.parquet', 256, 5, empty, empty, + SimpleStats.empty_stats(), SimpleStats.empty_stats(), 0, 4, 2, 1, []) + data = DataSplit([file], empty, 0, data_deletion_files=[DeletionFile('dv/one', 2, 10, 3)], + snapshot_id=7, bucket_path='bucket-0', total_buckets=1) + payload = IndexFileMeta('full-text' if full_text else 'ivf-flat', 'payload', 1024, 4, + external_path='s3://bucket/payload') + split_class = PrimaryKeyFullTextSearchSplit if full_text else PrimaryKeyVectorSearchSplit + source_class = PrimaryKeyFullTextScanPlan if full_text else PrimaryKeyVectorScanPlan + source_split = split_class( + data, (payload,), ('tail.parquet',), {'tail.parquet': (Range(1, 2), Range(4, 4))}) + source = source_class(7, [source_split], snapshot) + with patch.dict('sys.modules', {'pypaimon_rust': sdk, 'pypaimon_rust.datafusion': datafusion}): + native = _plan_to_native(pickle.loads(pickle.dumps(source)), full_text) + restored = _plan_from_native(search_table, native, full_text, True) + expected_type = datafusion.PrimaryKeyFullTextSearchSplit if full_text else datafusion.PrimaryKeyVectorSearchSplit + assert type(native.splits()[0]) is expected_type + assert type(native.splits()[0].data_split()) is datafusion.Split + assert serialize_split_v1(restored.splits()[0].data_split) == serialize_split_v1(data) + assert restored.snapshot() == snapshot and restored.snapshot_id == 7 + assert restored.splits()[0].payloads == (payload,) + assert restored.splits()[0].uncovered_data_files == ('tail.parquet',) + assert restored.splits()[0].row_ranges_by_file == source.splits()[0].row_ranges_by_file + + +@pytest.mark.parametrize('full_text', [False, True]) +@pytest.mark.parametrize('indexed', [False, True]) +def test_search_plan_conversion_rejects_splits_from_the_other_search_family(search_table, full_text, indexed): + from pypaimon.execution.native.read.native_search_plan import _plan_from_native, _plan_to_native + from pypaimon.table.source.full_text_search_split import IndexFullTextSearchSplit, RawFullTextSearchSplit + from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit, RawVectorSearchSplit + from pypaimon.utils.range import Range + + sdk, datafusion = _typed_native_sdk() + if full_text: + wrong = IndexVectorSearchSplit(0, 1, []) if indexed else RawVectorSearchSplit([Range(0, 1)]) + wrong_native = (datafusion.IndexVectorSearchSplit(0, 1, []) if indexed else + datafusion.RawVectorSearchSplit([(0, 1)])) + else: + wrong = IndexFullTextSearchSplit('body', 0, 1, []) if indexed else RawFullTextSearchSplit([Range(0, 1)]) + wrong_native = (datafusion.IndexFullTextSearchSplit('body', 0, 1, []) if indexed else + datafusion.RawFullTextSearchSplit([(0, 1)])) + source = (FullTextScanPlan if full_text else VectorSearchScanPlan)([wrong], None) + native = SimpleNamespace(snapshot=lambda: None, splits=lambda: [wrong_native]) + with patch.dict('sys.modules', {'pypaimon_rust': sdk, 'pypaimon_rust.datafusion': datafusion}): + with pytest.raises(TypeError, match='search split'): + _plan_to_native(source, full_text) + with pytest.raises(TypeError, match='search split'): + _plan_from_native(search_table, native, full_text, False) + + +@pytest.mark.parametrize('family', ['vector', 'batch', 'text']) +def test_selected_native_reader_does_not_probe_or_replay_an_older_sdk(search_table, family): + class OlderSdk(ModuleType): + def __getattr__(self, name): + if name in ('VectorScanPlan', 'FullTextScanPlan'): + raise AttributeError('required typed search plan API is missing') + raise AssertionError('Unexpected Native SDK capability probe: ' + name) + + sdk, datafusion = ModuleType('pypaimon_rust'), OlderSdk('pypaimon_rust.datafusion') + sdk.datafusion = datafusion + native, _, core_read = _native_builder(family) + builder = _builder(search_table, family) + _, read_factory = _factories(builder, family) + with patch.dict('sys.modules', {'pypaimon_rust': sdk, 'pypaimon_rust.datafusion': datafusion}), \ + _forbid_python_execution(), patch(_native_factory(family), return_value=native): + reader = read_factory() + with pytest.raises(AttributeError, match='required typed search plan API'): + (reader.read_batch_plan(_ordinary_plan(family)) if family == 'batch' else + reader.read_plan(_ordinary_plan(family))) + assert reader._python_executor is None + core_read.read_plan.assert_not_called() + core_read.read_batch_plan.assert_not_called() diff --git a/paimon-python/pypaimon/tests/native_search_plan_test.py b/paimon-python/pypaimon/tests/native_search_plan_test.py index 3c84d473d468..8f7cd86eb500 100644 --- a/paimon-python/pypaimon/tests/native_search_plan_test.py +++ b/paimon-python/pypaimon/tests/native_search_plan_test.py @@ -437,3 +437,28 @@ def test_pk_plan_retains_snapshot_and_sources_after_primary_key_update(source, p _assert_results(transported, expected) assert vectors._rows(table, actual[0]).to_pylist() == expected_rows assert vectors._rows(table, transported[0]).to_pylist() == expected_rows + + +@pytest.mark.parametrize('family', ['vector', 'batch', 'full-text']) +def test_public_sdk_plan_constructs_from_typed_sources_and_snapshot(table, family): + from pypaimon_rust import datafusion as sdk + from pypaimon.execution.native.read.native_search_plan import _plan_from_native, _plan_to_native + + builder = _builder(table, family, native=True) + with _native_only(): + source = _scan(builder, family).scan() + reader = _reader(builder, family) + expected = _read(reader, family, source) + full_text = family == 'full-text' + native = _plan_to_native(pickle.loads(pickle.dumps(source)), full_text) + cls = sdk.FullTextScanPlan if full_text else sdk.VectorScanPlan + reconstructed = cls(native.splits(), native.snapshot()) + assert reconstructed.snapshot().id() == source.snapshot().id + assert all(isinstance(split, (sdk.IndexFullTextSearchSplit, sdk.RawFullTextSearchSplit)) + if full_text else isinstance(split, (sdk.IndexVectorSearchSplit, sdk.RawVectorSearchSplit)) + for split in reconstructed.splits()) + assert not hasattr(reconstructed, 'metadata') and not hasattr(reconstructed, 'snapshot_id') + transported = _plan_from_native(table, reconstructed, full_text, False) + actual = _read(pickle.loads(pickle.dumps(reader)), family, pickle.loads(pickle.dumps(transported))) + assert pickle.dumps(transported.splits()) == pickle.dumps(source.splits()) + _assert_results(actual, expected) diff --git a/paimon-python/pypaimon/tests/search_task_execution_test.py b/paimon-python/pypaimon/tests/search_task_execution_test.py index 2029169fa665..1a6a327cce06 100644 --- a/paimon-python/pypaimon/tests/search_task_execution_test.py +++ b/paimon-python/pypaimon/tests/search_task_execution_test.py @@ -37,7 +37,7 @@ def _native_plan_conversion(): # These tests exercise callback ownership and cancellation. Source codec # construction has dedicated tests and must not require an installed SDK. with patch('pypaimon.execution.native.read.native_search_plan._plan_to_native', - side_effect=lambda plan, full_text=False, column=None: plan): + side_effect=lambda plan, full_text=False: plan): yield From 1b9cceeef73c63dd05c941139f11fe2101375dee Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Wed, 21 Oct 2026 23:03:47 +0800 Subject: [PATCH 4/9] Use unified native search execute entry points --- docs/docs/pypaimon/multimodal-search.md | 2 +- paimon-python/pypaimon/execution/README.md | 4 +- .../native/read/native_full_text_search.py | 2 +- .../native/read/native_hybrid_search.py | 2 +- .../native/read/native_vector_search.py | 4 +- .../tests/lazy_search_execution_test.py | 53 +++++++++++++++++++ .../tests/native_full_text_search_test.py | 11 ++-- .../tests/native_hybrid_search_test.py | 4 +- .../tests/native_index_result_scan_test.py | 10 ++-- .../tests/native_pk_full_text_search_test.py | 28 +++++----- .../tests/native_vector_context_test.py | 11 ++-- .../tests/native_vector_index_build_test.py | 4 +- .../tests/native_vector_search_test.py | 26 ++++----- 13 files changed, 107 insertions(+), 54 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-search.md b/docs/docs/pypaimon/multimodal-search.md index 0993d90521bb..21bd3c9bcfeb 100644 --- a/docs/docs/pypaimon/multimodal-search.md +++ b/docs/docs/pypaimon/multimodal-search.md @@ -380,7 +380,7 @@ and refinement; its internal shard protocol is separate from the public source plan. Hybrid builders expose `execute(executor=...)` for the same scheduler. -The existing `execute_local()` methods retain their local defaults. +PyPaimon's existing `execute_local()` methods retain their local defaults. ## Search Hybrid diff --git a/paimon-python/pypaimon/execution/README.md b/paimon-python/pypaimon/execution/README.md index 8e9a33138097..9dff4d14945b 100644 --- a/paimon-python/pypaimon/execution/README.md +++ b/paimon-python/pypaimon/execution/README.md @@ -164,7 +164,9 @@ Readers own the table, schema, query and filter. Native selection is retained after transport, and source subsets never add missing indexes or raw ranges. An absent snapshot describes an empty capture and never selects latest. Shard callbacks use an internal binding bridge; they are not part of Rust's public -Plan/Read API. +Plan/Read API. SDK vector, batch-vector, full-text and hybrid builders share +one `execute()` entry point. PyPaimon's existing public search methods delegate +to it through the Native adapters. ## Corresponding execution APIs diff --git a/paimon-python/pypaimon/execution/native/read/native_full_text_search.py b/paimon-python/pypaimon/execution/native/read/native_full_text_search.py index 9ba1505c6201..9ce5b5397b57 100644 --- a/paimon-python/pypaimon/execution/native/read/native_full_text_search.py +++ b/paimon-python/pypaimon/execution/native/read/native_full_text_search.py @@ -75,7 +75,7 @@ def try_native_full_text_search(builder): table = builder._table primary_key = builder._primary_key_full_text_definition() is not None # Once execution starts, propagate errors instead of running a second search. - result = native.execute_local() + result = native.execute() if primary_key: return PrimaryKeyScoredResult._from_native(table, result) return DictBasedScoredIndexResult(result.row_ids()) diff --git a/paimon-python/pypaimon/execution/native/read/native_hybrid_search.py b/paimon-python/pypaimon/execution/native/read/native_hybrid_search.py index d8ac3578299c..660c59593a5e 100644 --- a/paimon-python/pypaimon/execution/native/read/native_hybrid_search.py +++ b/paimon-python/pypaimon/execution/native/read/native_hybrid_search.py @@ -66,7 +66,7 @@ def try_native_hybrid_search(builder, executor=None): return None # Propagate execution failures: retrying Python can select a newer snapshot. if executor is None: - result = native.execute_local() + result = native.execute() else: from pypaimon.execution.native.read.search_execution import native_task_executor callback = native_task_executor(table, executor) diff --git a/paimon-python/pypaimon/execution/native/read/native_vector_search.py b/paimon-python/pypaimon/execution/native/read/native_vector_search.py index 83ea544fc5d6..3499eb5009ed 100644 --- a/paimon-python/pypaimon/execution/native/read/native_vector_search.py +++ b/paimon-python/pypaimon/execution/native/read/native_vector_search.py @@ -70,10 +70,10 @@ def try_native_vector_search(builder, queries, batch=False): # against a potentially newer snapshot. if batch: native.with_query_vectors(queries) - results = native.execute_batch_local() + results = native.execute() return [DictBasedScoredIndexResult(result.row_ids()) for result in results] native.with_query_vector(queries) - result = native.execute_local() + result = native.execute() if builder._is_primary_key_vector_search(): return PrimaryKeyScoredResult._from_native(table, result) return DictBasedScoredIndexResult(result.row_ids()) diff --git a/paimon-python/pypaimon/tests/lazy_search_execution_test.py b/paimon-python/pypaimon/tests/lazy_search_execution_test.py index 4e0b5b6799ff..527a8d99df4c 100644 --- a/paimon-python/pypaimon/tests/lazy_search_execution_test.py +++ b/paimon-python/pypaimon/tests/lazy_search_execution_test.py @@ -616,3 +616,56 @@ def __getattr__(self, name): assert reader._python_executor is None core_read.read_plan.assert_not_called() core_read.read_batch_plan.assert_not_called() + + +@pytest.mark.parametrize('family', ['vector', 'batch', 'text', 'hybrid']) +@pytest.mark.parametrize('failure', [False, True]) +def test_native_search_adapters_use_only_sdk_execute(search_table, family, failure): + from pypaimon.execution.native.read import native_full_text_search, native_hybrid_search, native_vector_search + + expected = SimpleNamespace(row_ids=lambda: {7: 0.75}) + error = OSError('Native execution started') + + class ExecuteOnlyBuilder: + def __init__(self): + self.execute = Mock(side_effect=error if failure else None, + return_value=[expected, expected] if family == 'batch' else expected) + for method in ('with_query_vector', 'with_query_vectors', 'with_limit', 'with_ranker', + 'add_vector_route', 'add_full_text_route'): + setattr(self, method, Mock(return_value=self)) + + def __getattr__(self, name): + if name in ('execute_local', 'execute_batch_local'): + raise AssertionError('Native adapter probed a removed SDK method: ' + name) + raise AttributeError(name) + + core = ExecuteOnlyBuilder() + with ExitStack() as stack: + if family in ('vector', 'batch'): + stack.enter_context(patch.object(native_vector_search, '_native_vector_builder', return_value=core)) + builder = _builder(search_table, family) + queries = builder._query_vectors if family == 'batch' else builder._query_vector + operation = lambda: native_vector_search.try_native_vector_search(builder, queries, family == 'batch') + elif family == 'text': + stack.enter_context(patch.object(native_full_text_search, '_native_full_text_builder', return_value=core)) + operation = lambda: native_full_text_search.try_native_full_text_search(_builder(search_table, 'text')) + else: + stack.enter_context(patch.object(native_hybrid_search, 'native_runtime_available', return_value=True)) + stack.enter_context(patch.object(native_hybrid_search, '_rest_catalog_supported', return_value=True)) + stack.enter_context(patch.object(native_hybrid_search, '_native_table', return_value=SimpleNamespace( + new_hybrid_search_builder=Mock(return_value=core)))) + builder = search_table.new_hybrid_search_builder().with_limit(2) + builder.add_vector_route('embedding', [1.0, 0.0], 2) + builder.add_full_text_route('body', '{"match":{"query":"needle"}}', 2) + operation = lambda: native_hybrid_search.try_native_hybrid_search(builder) + if failure: + with pytest.raises(OSError) as caught: + operation() + assert caught.value is error + else: + results = operation() + if family != 'batch': + results = [results] + assert [{row: result.score_getter()(row) for row in result.results()} for result in results] == ( + [{7: 0.75}] * (2 if family == 'batch' else 1)) + core.execute.assert_called_once_with() diff --git a/paimon-python/pypaimon/tests/native_full_text_search_test.py b/paimon-python/pypaimon/tests/native_full_text_search_test.py index 51d200b626b6..36cc51c186ef 100644 --- a/paimon-python/pypaimon/tests/native_full_text_search_test.py +++ b/paimon-python/pypaimon/tests/native_full_text_search_test.py @@ -276,24 +276,25 @@ def test_no_full_text_definition_is_empty_in_every_mode(rest_catalog, mode): assert indexes._ids(table) == list(range(8)) -def test_scan_read_plan_pins_snapshot_and_rejects_other_filters(rest_catalog): +def test_scan_read_plan_pins_snapshot_and_accepts_reader_filters(rest_catalog): catalog, _ = rest_catalog table = _table(catalog) native = (_native_table(table).new_full_text_search_builder() .with_query('content', match_query('paimon')).with_limit(64)) plan = native.new_full_text_scan().scan() before = table.snapshot_manager().get_latest_snapshot().id - assert plan.snapshot_id() == before + assert plan.snapshot().id() == before reader = native.new_full_text_read() indexes._append(table, _rows(8, 4), schema=SCHEMA) assert set(reader.read_plan(plan).row_ids()) == {0, 1, 2, 4, 5, 6, 7} - assert set(native.execute_local().row_ids()) == {0, 1, 2, 4, 5, 6, 7, 8, 9, 10, 11} + assert set(native.execute().row_ids()) == {0, 1, 2, 4, 5, 6, 7, 8, 9, 10, 11} predicate = table.new_read_builder().new_predicate_builder().equal('id', 0) changed = (_native_table(table).new_full_text_search_builder() .with_query('content', match_query('paimon')).with_limit(64)) changed.with_filter(_predicate_to_native(predicate)) - with pytest.raises(Exception, match='different table, column or filter'): - changed.new_full_text_read().read_plan(plan) + filtered = changed.new_full_text_read().read_plan(plan) + assert filtered.snapshot_id() == before + assert set(filtered.row_ids()) == {0} assert indexes._ids(table) == list(range(12)) diff --git a/paimon-python/pypaimon/tests/native_hybrid_search_test.py b/paimon-python/pypaimon/tests/native_hybrid_search_test.py index 3eed083379da..2ecd100e8168 100644 --- a/paimon-python/pypaimon/tests/native_hybrid_search_test.py +++ b/paimon-python/pypaimon/tests/native_hybrid_search_test.py @@ -239,9 +239,9 @@ def test_core_binding_rejects_invalid_configuration_before_empty_search(rest_cat table = indexes._create(catalog, schema=SCHEMA, partitioned=True, options={'read.native.enabled': 'true'}) native = _native_table(table) with pytest.raises(Exception, match='Routes cannot be empty'): - native.new_hybrid_search_builder().with_limit(2).execute_local() + native.new_hybrid_search_builder().with_limit(2).execute() with pytest.raises(Exception, match='Limit must be positive'): - native.new_hybrid_search_builder().add_vector_route('embedding', [0., 0.], 2).execute_local() + native.new_hybrid_search_builder().add_vector_route('embedding', [0., 0.], 2).execute() with pytest.raises(Exception, match='Unsupported hybrid ranker'): native.new_hybrid_search_builder().with_ranker('unknown') with pytest.raises(Exception, match='Limit must be positive'): diff --git a/paimon-python/pypaimon/tests/native_index_result_scan_test.py b/paimon-python/pypaimon/tests/native_index_result_scan_test.py index 33ede91810f4..ef646e8026fd 100644 --- a/paimon-python/pypaimon/tests/native_index_result_scan_test.py +++ b/paimon-python/pypaimon/tests/native_index_result_scan_test.py @@ -97,7 +97,7 @@ def test_core_search_result_passes_directly_to_scan_and_scored_splits_round_trip from pypaimon_rust.datafusion import GlobalIndexResult, Split table, _ = source native = text._native_table(table) - result = text._native(table).execute_local() + result = text._native(table).execute() assert isinstance(result, GlobalIndexResult) scan = native.new_read_builder().with_projection(['id']).new_scan().with_global_index_result(result) for plan in [scan.plan(), scan.plan()]: @@ -279,18 +279,18 @@ def test_all_core_searches_return_one_result_for_read_and_scan(rest_catalog, fam native = text._native_table(table) if family == 'vector': results = [(native.new_vector_search_builder().with_vector_column('embedding') - .with_query_vector([0., 0.]).with_limit(3).execute_local())] + .with_query_vector([0., 0.]).with_limit(3).execute())] elif family == 'batch-vector': results = (native.new_batch_vector_search_builder().with_vector_column('embedding') - .with_query_vectors([[0., 0.], [7., 0.]]).with_limit(3).execute_batch_local()) + .with_query_vectors([[0., 0.], [7., 0.]]).with_limit(3).execute()) assert len(results) == 2 elif family == 'full-text': - results = [text._native(table, 'paimon', 3).execute_local()] + results = [text._native(table, 'paimon', 3).execute()] else: results = [(native.new_hybrid_search_builder().with_limit(3) .add_vector_route('embedding', [0., 0.], 3, 1., {}) .add_full_text_route('content', hybrid.match_query('paimon'), 3, 1., {}) - .execute_local())] + .execute())] for result in results: assert isinstance(result, GlobalIndexResult) assert result.snapshot_id() == table.snapshot_manager().get_latest_snapshot().id diff --git a/paimon-python/pypaimon/tests/native_pk_full_text_search_test.py b/paimon-python/pypaimon/tests/native_pk_full_text_search_test.py index ccb2cda388bb..288600e147c7 100644 --- a/paimon-python/pypaimon/tests/native_pk_full_text_search_test.py +++ b/paimon-python/pypaimon/tests/native_pk_full_text_search_test.py @@ -97,7 +97,7 @@ def test_java_pk_archives_preserve_positions_scores_and_snapshot(source, snapsho @pytest.mark.parametrize('projection', [['id'], ['content', 'id'], []]) def test_core_result_keeps_scored_selections_and_materializes_only_the_projection(source, projection): table, _ = source - result = _native(table).execute_local() + result = _native(table).execute() assert result.snapshot_id() == table.snapshot_manager().get_latest_snapshot().id with pytest.raises(Exception, match='physical file positions'): result.row_ids() @@ -117,8 +117,8 @@ def test_captured_plan_does_not_resolve_latest_again_and_accepts_another_query(s table, server = source builder = _native(table) plan = builder.new_full_text_scan().scan() - snapshot_id = plan.snapshot_id() - expected = _native(table, 'Flink', 1).execute_local().positions() + snapshot_id = plan.snapshot().id() + expected = _native(table, 'Flink', 1).execute().positions() reader = _native(table, 'Flink', 1).new_full_text_read() unavailable = server._mock_response(ErrorResponse('TABLE', 'pk_text', 'snapshot unavailable', 503), 503) with patch.object(server, '_table_snapshot_handle', return_value=unavailable): @@ -130,7 +130,7 @@ def test_captured_plan_does_not_resolve_latest_again_and_accepts_another_query(s def test_search_returns_metadata_without_opening_user_parquet_files(source): table, server = source - expected = _native(table).execute_local().positions() + expected = _native(table).execute().positions() files = list((Path(server.data_path) / server.warehouse / 'default' / 'pk_text' / 'bucket-0').glob('*.parquet')) assert expected and files with ExitStack() as restore: @@ -138,24 +138,20 @@ def test_search_returns_metadata_without_opening_user_parquet_files(source): hidden = path.with_suffix('.held') path.rename(hidden) restore.callback(hidden.rename, path) - result = _native(table).execute_local() + result = _native(table).execute() assert result.positions() == expected with pytest.raises(Exception, match='NotFound|not found|No such file|does not exist'): result.new_read_builder().with_projection(['id']).read() -@pytest.mark.parametrize('case', ['table', 'filter', 'kind', 'column']) -def test_pk_plan_cannot_be_read_under_another_context(source, case): +@pytest.mark.parametrize('case', ['family', 'filter', 'kind', 'column']) +def test_pk_plan_validates_source_family_and_reader_requirements(source, case): table, _ = source plan = _native(table).new_full_text_scan().scan() - if case == 'table': - catalog = table.catalog_environment.catalog_loader.load() - schema = table.table_schema - catalog.create_table('default.other_text', Schema( - schema.fields, schema.partition_keys, schema.primary_keys, schema.options), False) - reader = _native(catalog.get_table('default.other_text')).new_full_text_read() - with pytest.raises(Exception, match='different table, column or filter'): - reader.read_plan(plan) + if case == 'family': + from pypaimon_rust.datafusion import FullTextScanPlan, RawVectorSearchSplit + with pytest.raises(TypeError): + FullTextScanPlan([RawVectorSearchSplit([(0, 1)])], plan.snapshot()) elif case == 'filter': reader = _native(table) reader.with_filter(_predicate_to_native(table.new_read_builder().new_predicate_builder().equal('id', 1))) @@ -164,7 +160,7 @@ def test_pk_plan_cannot_be_read_under_another_context(source, case): elif case == 'kind': changed = table.copy({'pk-full-text.index.columns': ''}) reader = _native(changed).new_full_text_read() - with pytest.raises(Exception, match='different table, column or filter'): + with pytest.raises(Exception, match='Global full-text read requires global row-ID sources'): reader.read_plan(plan) else: reader = (_native_table(table).new_full_text_search_builder() diff --git a/paimon-python/pypaimon/tests/native_vector_context_test.py b/paimon-python/pypaimon/tests/native_vector_context_test.py index 8163083b5555..9c4608a36fde 100644 --- a/paimon-python/pypaimon/tests/native_vector_context_test.py +++ b/paimon-python/pypaimon/tests/native_vector_context_test.py @@ -49,8 +49,8 @@ def test_unused_context_attributes_preserve_native_vector_execution(option, batc native_table.new_vector_search_builder.return_value = core native_table.new_batch_vector_search_builder.return_value = core expected = [{7: -0.25, 9: -1.5}, {11: -2.0}] - core.execute_local.return_value.row_ids.return_value = expected[0] - core.execute_batch_local.return_value = [Mock(row_ids=Mock(return_value=scores)) for scores in expected] + core.execute.return_value = ([Mock(row_ids=Mock(return_value=scores)) for scores in expected] + if batch else Mock(row_ids=Mock(return_value=expected[0]))) queries = [[1.0, 0.0], [0.0, 1.0]] if batch else [1.0, 0.0] with patch.object(loader, 'load', side_effect=AssertionError('Python catalog was reconstructed')), \ patch.object(native_vector_search, 'native_runtime_available', return_value=True), \ @@ -66,13 +66,12 @@ def test_unused_context_attributes_preserve_native_vector_execution(option, batc unused.assert_not_called() if batch: core.with_query_vectors.assert_called_once_with(queries) - core.execute_batch_local.assert_called_once_with() - core.execute_local.assert_not_called() results = result else: core.with_query_vector.assert_called_once_with(queries) - core.execute_local.assert_called_once_with() - core.execute_batch_local.assert_not_called() results = [result] expected = expected[:1] + core.execute.assert_called_once_with() + core.execute_local.assert_not_called() + core.execute_batch_local.assert_not_called() assert [{row_id: item.score_getter()(row_id) for row_id in item.results()} for item in results] == expected diff --git a/paimon-python/pypaimon/tests/native_vector_index_build_test.py b/paimon-python/pypaimon/tests/native_vector_index_build_test.py index 38b56791a6e2..04fd3e5d9803 100644 --- a/paimon-python/pypaimon/tests/native_vector_index_build_test.py +++ b/paimon-python/pypaimon/tests/native_vector_index_build_test.py @@ -81,7 +81,7 @@ def _assert_ids(table, expected, kind): .with_query_vector([1., 1.]).with_limit(64).with_options(search_options).execute_local()) assert sorted(result.results()) == expected result = (_native_table(table).new_vector_search_builder().with_vector_column('embedding') - .with_query_vector([1., 1.]).with_limit(64).with_options(search_options).execute_local()) + .with_query_vector([1., 1.]).with_limit(64).with_options(search_options).execute()) assert sorted(result.row_ids()) == expected @@ -135,7 +135,7 @@ def test_native_search_filters_live_null_vectors_with_deletion_vectors(rest_cata indexes._commit(table, _build(table, kind, generic._options(kind))) options = {'diskann.l_search': '64'} if kind == 'diskann' else {'ivf.nprobe': '64'} result = (_native_table(table).new_vector_search_builder().with_vector_column('embedding') - .with_query_vector([1., 1.]).with_limit(64).with_options(options).execute_local()) + .with_query_vector([1., 1.]).with_limit(64).with_options(options).execute()) assert sorted(result.row_ids()) == [1, 4] # Python table-level planning also uses the core low-level index reader, # avoiding the third-party scalar FFI's fixed-result-length restriction. diff --git a/paimon-python/pypaimon/tests/native_vector_search_test.py b/paimon-python/pypaimon/tests/native_vector_search_test.py index af8db8dc3975..b7b8c54007e7 100644 --- a/paimon-python/pypaimon/tests/native_vector_search_test.py +++ b/paimon-python/pypaimon/tests/native_vector_search_test.py @@ -157,7 +157,7 @@ def test_scan_plan_is_reusable_and_snapshot_pinned(vector_table): builder = table.new_vector_search_builder().with_vector_column('embedding') scan = builder.new_vector_search_scan() plan = scan.scan() - snapshot = plan.snapshot_id() + snapshot = plan.snapshot().id() _write(vector_table, [{'id': 7, 'embedding': [1, 0], 'pt': 0}]) builder.with_limit(10).with_query_vector([1, 0]) reader = builder.new_vector_search_read() @@ -165,12 +165,14 @@ def test_scan_plan_is_reusable_and_snapshot_pinned(vector_table): assert result.snapshot_id() == snapshot assert len(result) == 5 assert reader.read_plan(plan).row_ids() == result.row_ids() - assert len(builder.execute_local()) == 6 + assert len(builder.execute()) == 6 filtered = table.new_vector_search_builder().with_vector_column('embedding').with_query_vector([1, 0]).with_limit(2) predicate = vector_table.new_read_builder().new_predicate_builder().equal('id', 1) filtered.with_filter(_predicate_to_native(predicate)) - with pytest.raises(Exception, match='plan.*reader|context|different'): - filtered.new_vector_search_read().read_plan(plan) + selected = filtered.new_vector_search_read().read_plan(plan) + assert selected.snapshot_id() == snapshot + assert [value for batch in selected.new_read_builder().with_projection(['id']).read() + for value in batch.column('id').to_pylist()] == [1] @pytest.fixture @@ -308,7 +310,7 @@ def test_mixed_index_metrics_are_rejected(vector_table, filtered, mode): if filtered: direct.with_filter(_predicate_to_native(predicate)) with pytest.raises(Exception, match='Cannot merge vector indexes with different metrics'): - direct.execute_local() + direct.execute() builders = [_single(table), table.new_batch_vector_search_builder().with_vector_column('embedding') .with_query_vectors([[1, 0], [0, 1]]).with_limit(3)] for builder in builders: @@ -389,7 +391,7 @@ def configure(builder, direct=False): assert all(len(scores) == 3 for scores in expected) direct_table = _native_table(native) direct = configure(direct_table.new_batch_vector_search_builder().with_vector_column('embedding') - .with_query_vectors(queries).with_limit(3), direct=True).execute_batch_local() + .with_query_vectors(queries).with_limit(3), direct=True).execute() for result, scores in zip(direct, expected): _assert_scores(result.row_ids(), scores) with patch.object(DataEvolutionVectorRead, 'read_plan', side_effect=AssertionError('Python search ran')), \ @@ -453,7 +455,7 @@ def configure(builder, direct=False): direct_table = _native_table(native) for query, scores in zip(queries, expected): result = configure(direct_table.new_vector_search_builder().with_vector_column('embedding') - .with_query_vector(vector=query).with_limit(1), direct=True).execute_local() + .with_query_vector(vector=query).with_limit(1), direct=True).execute() _assert_scores(result.row_ids(), scores) with patch.object(DataEvolutionVectorRead, 'read_plan', side_effect=AssertionError('Python search ran')), \ patch.object(BatchVectorSearchReadImpl, 'read_batch_plan', side_effect=AssertionError('Python search ran')): @@ -488,7 +490,7 @@ def configure(builder, direct=False): expected = _scores(configure(_single(classic)).execute_local()) assert len(expected) == 1 direct = configure(_native_table(native).new_vector_search_builder().with_vector_column('embedding') - .with_query_vector(vector=[1, 0]).with_limit(3), direct=True).execute_local() + .with_query_vector(vector=[1, 0]).with_limit(3), direct=True).execute() _assert_scores(direct.row_ids(), expected) with patch.object(DataEvolutionVectorRead, 'read_plan', side_effect=AssertionError('Python search ran')), \ patch.object(BatchVectorSearchReadImpl, 'read_batch_plan', side_effect=AssertionError('Python search ran')): @@ -548,7 +550,7 @@ def test_unsupported_conjunct_on_same_indexed_field_is_not_exact(rest_catalog, m for predicate, count in [(is_null, 1), (impossible, 0)]: direct = (direct_table.new_vector_search_builder().with_vector_column('embedding') .with_query_vector(vector=[1, 0]).with_limit(2) - .with_filter(_predicate_to_native(predicate)).execute_local()) + .with_filter(_predicate_to_native(predicate)).execute()) assert len(direct) == count with patch.object(DataEvolutionVectorRead, 'read_plan', side_effect=AssertionError('Python search ran')), \ patch.object(BatchVectorSearchReadImpl, 'read_batch_plan', @@ -602,10 +604,10 @@ def configure(builder, direct=False): direct_table = _native_table(native) for query, scores in zip(queries, expected): result = configure(direct_table.new_vector_search_builder().with_vector_column('embedding') - .with_query_vector(vector=query).with_limit(10), direct=True).execute_local() + .with_query_vector(vector=query).with_limit(10), direct=True).execute() _assert_scores(result.row_ids(), scores) direct_batch = configure(direct_table.new_batch_vector_search_builder().with_vector_column('embedding') - .with_query_vectors(queries).with_limit(10), direct=True).execute_batch_local() + .with_query_vectors(queries).with_limit(10), direct=True).execute() for result, scores in zip(direct_batch, expected): _assert_scores(result.row_ids(), scores) with patch.object(DataEvolutionVectorRead, 'read_plan', side_effect=AssertionError('Python search ran')), \ @@ -674,7 +676,7 @@ def configure(builder, direct=False): assert [len(scores) for scores in expected] == [2, 2] direct_table = _native_table(native) direct = configure(direct_table.new_batch_vector_search_builder().with_vector_column('embedding') - .with_query_vectors(queries).with_limit(10), direct=True).execute_batch_local() + .with_query_vectors(queries).with_limit(10), direct=True).execute() for result, scores in zip(direct, expected): _assert_scores(result.row_ids(), scores) with patch.object(DataEvolutionVectorRead, 'read_plan', side_effect=AssertionError('Python search ran')), \ From bd2126200a6c4c2eeaa0781f43f919db304cf932 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Wed, 21 Oct 2026 23:21:50 +0800 Subject: [PATCH 5/9] Fix transported search sources and scheduled failure handling --- docs/docs/pypaimon/multimodal-search.md | 4 +- paimon-python/pypaimon/execution/README.md | 4 +- .../native/read/native_search_plan.py | 10 +- .../search/distributed_full_text_search.py | 12 ++- .../python/read/search/distributed_search.py | 16 +++ .../read/search/distributed_vector_search.py | 12 +-- .../read/search/primary_key_full_text_scan.py | 17 +-- .../read/search/primary_key_vector_scan.py | 17 ++- .../tests/lazy_search_execution_test.py | 11 +- .../pypaimon/tests/native_search_plan_test.py | 6 +- .../primary_key_index_definitions_test.py | 31 ++++++ .../tests/search_task_execution_test.py | 101 ++++++++++++++++++ 12 files changed, 207 insertions(+), 34 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-search.md b/docs/docs/pypaimon/multimodal-search.md index 21bd3c9bcfeb..57513fae0f79 100644 --- a/docs/docs/pypaimon/multimodal-search.md +++ b/docs/docs/pypaimon/multimodal-search.md @@ -373,8 +373,8 @@ results = builder.new_batch_vector_search_read().read_batch_plan(selected) ``` The supplied splits are authoritative: readers do not add uncovered raw ranges -or rediscover other search indexes. A plan with no captured snapshot stays -empty rather than selecting the latest snapshot. Query vectors, limits and +or rediscover other search indexes. A Native plan without a captured snapshot +must have no sources and never selects the latest snapshot. Query vectors, limits and filters belong to the reader. Rust Core coordinates global candidate selection and refinement; its internal shard protocol is separate from the public source plan. diff --git a/paimon-python/pypaimon/execution/README.md b/paimon-python/pypaimon/execution/README.md index 9dff4d14945b..699c9b3dad01 100644 --- a/paimon-python/pypaimon/execution/README.md +++ b/paimon-python/pypaimon/execution/README.md @@ -162,8 +162,8 @@ made by Python scans or sent to another process. Existing snapshot and index metadata codecs are private adapters, rather than a public JSON plan protocol. Readers own the table, schema, query and filter. Native selection is retained after transport, and source subsets never add missing indexes or raw ranges. -An absent snapshot describes an empty capture and never selects latest. Shard -callbacks use an internal binding bridge; they are not part of Rust's public +For Native, a plan without a captured snapshot must have no sources and never +selects latest. Shard callbacks use an internal binding bridge; they are not part of Rust's public Plan/Read API. SDK vector, batch-vector, full-text and hybrid builders share one `execute()` entry point. PyPaimon's existing public search methods delegate to it through the Native adapters. diff --git a/paimon-python/pypaimon/execution/native/read/native_search_plan.py b/paimon-python/pypaimon/execution/native/read/native_search_plan.py index 95cebb37878c..e40992087844 100644 --- a/paimon-python/pypaimon/execution/native/read/native_search_plan.py +++ b/paimon-python/pypaimon/execution/native/read/native_search_plan.py @@ -80,7 +80,7 @@ def _plan_from_native(table, native, full_text, primary_key): return (FullTextScanPlan if full_text else VectorSearchScanPlan)(splits, snapshot) -def _split_to_native(split, full_text): +def _split_to_native(split, full_text, snapshot=None): from pypaimon_rust import datafusion as sdk from pypaimon.table.source.full_text_search_split import IndexFullTextSearchSplit, RawFullTextSearchSplit from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit, RawVectorSearchSplit @@ -99,7 +99,11 @@ def ranges(values): if isinstance(split, (PrimaryKeyVectorSearchSplit, PrimaryKeyFullTextSearchSplit)): cls = sdk.PrimaryKeyFullTextSearchSplit if full_text else sdk.PrimaryKeyVectorSearchSplit - return cls(sdk.Split.deserialize(serialize_split_v1(split.data_split)), files(split.payloads), + data_split = split.data_split + if data_split.snapshot_id is None and snapshot is not None: + data_split = copy(data_split) + data_split.snapshot_id = snapshot.id + return cls(sdk.Split.deserialize(serialize_split_v1(data_split)), files(split.payloads), list(split.uncovered_data_files), {name: ranges(values) for name, values in split.row_ranges_by_file.items()}) if isinstance(split, IndexFullTextSearchSplit): @@ -119,7 +123,7 @@ def _plan_to_native(plan, full_text): captured = plan.snapshot() snapshot = None if captured is None else sdk.Snapshot._from_json(JSON.to_json(captured)) cls = sdk.FullTextScanPlan if full_text else sdk.VectorScanPlan - return cls([_split_to_native(split, full_text) for split in plan.splits()], snapshot) + return cls([_split_to_native(split, full_text, captured) for split in plan.splits()], snapshot) def _capture_builder(builder): diff --git a/paimon-python/pypaimon/execution/python/read/search/distributed_full_text_search.py b/paimon-python/pypaimon/execution/python/read/search/distributed_full_text_search.py index fa1497e60cc8..010a609fd4f9 100644 --- a/paimon-python/pypaimon/execution/python/read/search/distributed_full_text_search.py +++ b/paimon-python/pypaimon/execution/python/read/search/distributed_full_text_search.py @@ -17,14 +17,13 @@ """Python full-text tasks preserve corpus statistics and plan-ordered merging.""" -from contextlib import closing - from pypaimon.globalindex.global_index_result import GlobalIndexResult from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult from pypaimon.execution.python.read.search.full_text_read import ( DataEvolutionFullTextRead, _index_candidate_limit, _index_type, _raw_row_ranges, _split_search_splits, ) from pypaimon.execution.python.read.search import global_index_live_row_filter +from pypaimon.execution.python.read.search.distributed_search import closing_task_results from pypaimon.execution.python.read.search.distributed_vector_search import _index_tasks, _scores from pypaimon.table.source.search_execution import SearchTask from pypaimon.utils.range import Range @@ -47,7 +46,10 @@ def _worker_reader(self): def read(self, splits): # Keep bitmap construction and candidate data reads off the driver. index_splits, raw_splits = _split_search_splits(splits) - indexed = self._eval_column_query({self._text_columns[0].name: index_splits}, None) + splits_by_column = {} + for split in index_splits: + splits_by_column.setdefault(split.column_name, []).append(split) + indexed = self._eval_column_query(splits_by_column, None) raw = self._read_raw_search(_raw_row_ranges(raw_splits), _index_type(index_splits)) return indexed.or_(raw).top_k(self._limit) @@ -59,7 +61,7 @@ def _eval_column_query(self, splits_by_column, live_rows): reader = self._worker_reader() reader._limit = _index_candidate_limit(items, self._limit) work = _index_tasks(_search_index, reader, items, self._filter, _prepare_scalar_filter) - with closing(self._executor.map(work, ordered=True)) as tasks: + with closing_task_results(self._executor.map(work, ordered=True)) as tasks: for _, scores in tasks: # Duplicate row IDs keep the first planned shard's score. for row_id, score in scores.items(): @@ -73,7 +75,7 @@ def _read_raw_search(self, ranges, index_type, include_row_ids=None): # Splitting this corpus would change document frequencies and BM25. # Keep the existing local corpus intact, but read/build/search on a worker. work = [SearchTask(_search_raw, self._worker_reader(), (ranges, index_type, include))] - with closing(self._executor.map(work)) as tasks: + with closing_task_results(self._executor.map(work)) as tasks: for _, scores in tasks: return DictBasedScoredIndexResult(scores) 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 d8af398a8c68..783798d3584a 100644 --- a/paimon-python/pypaimon/execution/python/read/search/distributed_search.py +++ b/paimon-python/pypaimon/execution/python/read/search/distributed_search.py @@ -17,9 +17,25 @@ """Attach distributed task execution to Python search readers.""" +from contextlib import contextmanager import math +@contextmanager +def closing_task_results(results): + """Close scheduled work without replacing an existing execution failure.""" + try: + yield results + except BaseException: + try: + results.close() + except BaseException: + pass + raise + else: + results.close() + + def create_distributed_reader(reader, executor): executor = executor.with_concurrency(executor.concurrency) from pypaimon.execution.python.read.search.vector_search_read import ( diff --git a/paimon-python/pypaimon/execution/python/read/search/distributed_vector_search.py b/paimon-python/pypaimon/execution/python/read/search/distributed_vector_search.py index 95618ee8eb7f..78ac0462cc16 100644 --- a/paimon-python/pypaimon/execution/python/read/search/distributed_vector_search.py +++ b/paimon-python/pypaimon/execution/python/read/search/distributed_vector_search.py @@ -18,7 +18,6 @@ """Python search tasks and global result coordination for distributed executors.""" from collections import Counter -from contextlib import closing import math from pypaimon.globalindex.batch_vector_search import BatchVectorSearch @@ -29,6 +28,7 @@ _filtered_raw_row_ranges, _offer_score, _scored_result, _split_search_splits, ) from pypaimon.execution.python.read.search import global_index_live_row_filter +from pypaimon.execution.python.read.search.distributed_search import closing_task_results from pypaimon.table.source.search_execution import SearchTask from pypaimon.utils.range import Range from pypaimon.utils.roaring_bitmap import RoaringBitmap64 @@ -55,7 +55,7 @@ def _search_index_splits(self, splits, query, search_limit, snapshot, batch=Fals context = (self._table, self._vector_column, query, search_limit, self._options, self._filter, self._partition_filter, snapshot) work = _index_tasks(_search_index_split, context, splits, self._filter) - with closing(self._executor.map(work, ordered=True)) as tasks: + with closing_task_results(self._executor.map(work, ordered=True)) as tasks: for _, (metric, scores) in tasks: if metric is not None: self._set_index_metric(metric) @@ -81,7 +81,7 @@ def _read_raw_search(self, raw_row_ranges, pre_filter, query_vector, # Keep each planned data split intact: its files may be vertical column # groups that must be read together to reconstruct a logical row. work = [SearchTask(_search_raw_split, context, split) for split in splits] - with closing(self._executor.map(work)) as tasks: + with closing_task_results(self._executor.map(work)) as tasks: for _, scores in tasks: for row_id, score in scores.items(): _offer_score(heap, self._limit, row_id, score) @@ -107,7 +107,7 @@ def _search_index_splits(self, splits, query, search_limit, snapshot, batch=Fals context = (self._table, self._vector_column, query, search_limit, self._options, self._filter, self._partition_filter, snapshot) work = _index_tasks(_search_batch_index_split, context, splits, self._filter) - with closing(self._executor.map(work, ordered=True)) as tasks: + with closing_task_results(self._executor.map(work, ordered=True)) as tasks: for _, (metric, scores) in tasks: if metric is not None: self._set_index_metric(metric) @@ -124,7 +124,7 @@ def _read_raw_batch_search(self, raw_row_ranges, pre_filter, index_type=None, sn context = (table_read, self._vector_column, self._query_vectors, self._limit, self._search_metric(index_type)) work = [SearchTask(_search_batch_raw_split, context, split) for split in splits] - with closing(self._executor.map(work)) as tasks: + with closing_task_results(self._executor.map(work)) as tasks: for _, results in tasks: for heap, scores in zip(heaps, results): for row_id, score in scores.items(): @@ -144,7 +144,7 @@ def _stream_rerank_candidates(self, candidates, union_candidates, query_vectors, self._search_metric(index_type), queries_by_row) heaps = [[] for _ in query_vectors] work = [SearchTask(_search_batch_refine_split, context, split) for split in splits] - with closing(self._executor.map(work)) as tasks: + with closing_task_results(self._executor.map(work)) as tasks: for _, results in tasks: for heap, scores in zip(heaps, results): for row_id, score in scores.items(): diff --git a/paimon-python/pypaimon/execution/python/read/search/primary_key_full_text_scan.py b/paimon-python/pypaimon/execution/python/read/search/primary_key_full_text_scan.py index b66e29c90283..a08ba3bb105c 100644 --- a/paimon-python/pypaimon/execution/python/read/search/primary_key_full_text_scan.py +++ b/paimon-python/pypaimon/execution/python/read/search/primary_key_full_text_scan.py @@ -16,6 +16,8 @@ # under the License. +from copy import copy + from pypaimon.table.source.search_types import (PrimaryKeyFullTextSearchSplit, PrimaryKeyFullTextScanPlan) from pypaimon.globalindex.indexed_split import IndexedSplit @@ -61,7 +63,7 @@ def scan(self): entry, self._definition.field_id, self._definition.index_type, self._partition_filter)) return PrimaryKeyFullTextScanPlan( - snapshot.id, _full_text_bucket_splits(source_splits, entries), snapshot) + snapshot.id, _full_text_bucket_splits(source_splits, entries, snapshot.id), snapshot) def _matches(entry, field_id, index_type, partition_filter): @@ -77,7 +79,7 @@ def _partition_key(partition): return tuple(partition.values) -def _full_text_bucket_splits(source_splits, entries): +def _full_text_bucket_splits(source_splits, entries, snapshot_id=None): payloads_by_bucket = {} for entry in entries: key = (_partition_key(entry.partition), entry.bucket) @@ -89,7 +91,7 @@ def _full_text_bucket_splits(source_splits, entries): continue key = (_partition_key(split.partition), split.bucket) bucket = buckets.setdefault(key, { - "partition": split.partition, "files": [], + "source": split, "files": [], "deletions": [], "names": set(), "has_deletions": False}) deletions = split.data_deletion_files if deletions is not None and len(deletions) != len(split.files): @@ -114,9 +116,12 @@ def _full_text_bucket_splits(source_splits, entries): continue payloads, covered = _current_payloads( bucket["files"], payloads_by_bucket.get(key, [])) - data_split = DataSplit( - bucket["files"], bucket["partition"], key[1], False, - bucket["deletions"] if bucket["has_deletions"] else None) + data_split = copy(bucket["source"]) + data_split._files = bucket["files"] + data_split.raw_convertible = False + data_split.data_deletion_files = bucket["deletions"] if bucket["has_deletions"] else None + if snapshot_id is not None: + data_split.snapshot_id = snapshot_id result.append(PrimaryKeyFullTextSearchSplit( data_split, tuple(payloads), tuple(data_file.file_name for data_file in bucket["files"] diff --git a/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_scan.py b/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_scan.py index 7d2c0ea334f5..ffe67e5fc0eb 100644 --- a/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_scan.py +++ b/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_scan.py @@ -16,6 +16,8 @@ # under the License. +from copy import copy + from pypaimon.table.source.search_types import (PrimaryKeyVectorSearchSplit, PrimaryKeyVectorScanPlan) from pypaimon.common.options.core_options import CoreOptions @@ -87,7 +89,7 @@ def scan(self): lambda entry: _matches(entry, self._vector_column.id, index_type, self._partition_filter)) return PrimaryKeyVectorScanPlan( - snapshot.id, _bucket_splits(source_splits, entries), snapshot) + snapshot.id, _bucket_splits(source_splits, entries, snapshot.id), snapshot) def _matches(entry, field_id, index_type, partition_filter): @@ -98,7 +100,7 @@ def _matches(entry, field_id, index_type, partition_filter): and (partition_filter is None or partition_filter.test(entry.partition))) -def _bucket_splits(source_splits, entries): +def _bucket_splits(source_splits, entries, snapshot_id=None): payloads_by_bucket = {} for entry in entries: key = (_partition_key(entry.partition), entry.bucket) @@ -110,7 +112,9 @@ def _bucket_splits(source_splits, entries): ranges_by_bucket.setdefault(key, {}).update(ranges) previous = combined.get(key) if previous is None: - combined[key] = split + combined[key] = copy(split) + if snapshot_id is not None: + combined[key].snapshot_id = snapshot_id continue files = list(previous.files) + list(split.files) deletions = None @@ -118,8 +122,11 @@ def _bucket_splits(source_splits, entries): or split.data_deletion_files is not None): deletions = list(previous.data_deletion_files or [None] * len(previous.files)) deletions.extend(split.data_deletion_files or [None] * len(split.files)) - combined[key] = DataSplit(files, previous.partition, previous.bucket, - False, deletions) + merged = copy(previous) + merged._files = files + merged.raw_convertible = False + merged.data_deletion_files = deletions + combined[key] = merged result = [] for split in combined.values(): diff --git a/paimon-python/pypaimon/tests/lazy_search_execution_test.py b/paimon-python/pypaimon/tests/lazy_search_execution_test.py index 527a8d99df4c..8b91fbb14426 100644 --- a/paimon-python/pypaimon/tests/lazy_search_execution_test.py +++ b/paimon-python/pypaimon/tests/lazy_search_execution_test.py @@ -526,7 +526,8 @@ def test_missing_primary_key_plan_snapshot_does_not_fall_back_to_latest(full_tex @pytest.mark.parametrize('full_text', [False, True]) -def test_primary_key_sources_round_trip_as_typed_splits(search_table, full_text): +@pytest.mark.parametrize('split_snapshot', [None, 7]) +def test_primary_key_sources_round_trip_as_typed_splits(search_table, full_text, split_snapshot): from pypaimon.execution.native.read.native_search_plan import _plan_from_native, _plan_to_native from pypaimon.index.index_file_meta import IndexFileMeta from pypaimon.manifest.schema.data_file_meta import DataFileMeta @@ -547,7 +548,7 @@ def test_primary_key_sources_round_trip_as_typed_splits(search_table, full_text) file = DataFileMeta.create('tail.parquet', 256, 5, empty, empty, SimpleStats.empty_stats(), SimpleStats.empty_stats(), 0, 4, 2, 1, []) data = DataSplit([file], empty, 0, data_deletion_files=[DeletionFile('dv/one', 2, 10, 3)], - snapshot_id=7, bucket_path='bucket-0', total_buckets=1) + snapshot_id=split_snapshot, bucket_path='bucket-0', total_buckets=1) payload = IndexFileMeta('full-text' if full_text else 'ivf-flat', 'payload', 1024, 4, external_path='s3://bucket/payload') split_class = PrimaryKeyFullTextSearchSplit if full_text else PrimaryKeyVectorSearchSplit @@ -561,7 +562,11 @@ def test_primary_key_sources_round_trip_as_typed_splits(search_table, full_text) expected_type = datafusion.PrimaryKeyFullTextSearchSplit if full_text else datafusion.PrimaryKeyVectorSearchSplit assert type(native.splits()[0]) is expected_type assert type(native.splits()[0].data_split()) is datafusion.Split - assert serialize_split_v1(restored.splits()[0].data_split) == serialize_split_v1(data) + from copy import copy + normalized = copy(data) + normalized.snapshot_id = 7 + assert serialize_split_v1(restored.splits()[0].data_split) == serialize_split_v1(normalized) + assert data.snapshot_id == split_snapshot assert restored.snapshot() == snapshot and restored.snapshot_id == 7 assert restored.splits()[0].payloads == (payload,) assert restored.splits()[0].uncovered_data_files == ('tail.parquet',) diff --git a/paimon-python/pypaimon/tests/native_search_plan_test.py b/paimon-python/pypaimon/tests/native_search_plan_test.py index 8f7cd86eb500..2a93bb5565d6 100644 --- a/paimon-python/pypaimon/tests/native_search_plan_test.py +++ b/paimon-python/pypaimon/tests/native_search_plan_test.py @@ -379,7 +379,7 @@ def map(self, tasks, ordered=False): @pytest.mark.parametrize('family', ['vector', 'full-text']) @pytest.mark.parametrize('filtered', [False, True]) -def test_pk_native_plans_export_physical_sources_for_python_readers(source, pk_vector_table, family, filtered): +def test_pk_plans_round_trip_between_backends(source, pk_vector_table, family, filtered): table = source[0] if family == 'full-text' else pk_vector_table def builder(native): @@ -390,9 +390,11 @@ def builder(native): return query classic = builder(False) - expected = _read(_reader(classic, family), family, _scan(classic, family).scan()) + python_plan = _scan(classic, family).scan() + expected = _read(_reader(classic, family), family, python_plan) native = builder(True) with _native_only(): + _assert_results(_read(_reader(native, family), family, pickle.loads(pickle.dumps(python_plan))), expected) plan = _scan(native, family).scan() actual = _read(_reader(native, family), family, plan) assert plan.snapshot_id == expected[0].snapshot_id diff --git a/paimon-python/pypaimon/tests/primary_key_index_definitions_test.py b/paimon-python/pypaimon/tests/primary_key_index_definitions_test.py index b3a5a1847071..12ad4339ca78 100644 --- a/paimon-python/pypaimon/tests/primary_key_index_definitions_test.py +++ b/paimon-python/pypaimon/tests/primary_key_index_definitions_test.py @@ -53,6 +53,37 @@ class PrimaryKeyIndexDefinitionsTest(unittest.TestCase): + def test_bucket_planners_preserve_capture_and_physical_metadata(self): + from pypaimon.execution.python.read.search.primary_key_full_text_scan import _full_text_bucket_splits + from pypaimon.table.source.deletion_file import DeletionFile + + partition = GenericRow([], []) + files = [DataFileMeta( + file_name=name, file_size=10, row_count=2, + min_key=None, max_key=None, key_stats=None, value_stats=None, + min_sequence_number=0, max_sequence_number=0, + schema_id=0, level=1, extra_files=[], file_source=1) + for name in ('a.parquet', 'b.parquet')] + deletion = DeletionFile('dv/index', 2, 10, 1) + for snapshot_id in (None, 7): + sources = [DataSplit( + [file], partition, 3, snapshot_id=snapshot_id, + data_deletion_files=[deletion] if ordinal == 0 else None, + bucket_path='s3://warehouse/table/bucket-3', total_buckets=4) + for ordinal, file in enumerate(files)] + full_text = _full_text_bucket_splits(sources, [], 7) + vector = _bucket_splits([(source, {}) for source in sources], [], 7) + for planned in (full_text, vector): + self.assertEqual(1, len(planned)) + data = planned[0].data_split + self.assertEqual(7, data.snapshot_id) + self.assertEqual('s3://warehouse/table/bucket-3', data.bucket_path) + self.assertEqual(4, data.total_buckets) + self.assertEqual(files, data.files) + self.assertEqual([deletion, None], data.data_deletion_files) + self.assertEqual([snapshot_id, snapshot_id], [source.snapshot_id for source in sources]) + self.assertEqual([[files[0]], [files[1]]], [source.files for source in sources]) + def test_definitions_follow_schema_order(self): schema = TableSchema( fields=[DataField(0, "id", AtomicType("INT")), diff --git a/paimon-python/pypaimon/tests/search_task_execution_test.py b/paimon-python/pypaimon/tests/search_task_execution_test.py index 1a6a327cce06..81721a955220 100644 --- a/paimon-python/pypaimon/tests/search_task_execution_test.py +++ b/paimon-python/pypaimon/tests/search_task_execution_test.py @@ -225,3 +225,104 @@ def close(self): assert next(results) == b'result' results.close() callback.raise_if_failed() + + +@pytest.mark.parametrize('columns,expected', [ + (['title'], []), (['body'], [1]), (['title', 'body'], [1]), +]) +def test_scheduled_full_text_read_keeps_each_source_column(columns, expected): + from pypaimon.common.options.core_options import CoreOptions + from pypaimon.common.options.options import Options + from pypaimon.execution.python.read.search.full_text_read import DataEvolutionFullTextRead + from pypaimon.globalindex.global_index_reader import _completed_future + from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult + from pypaimon.table.source.full_text_search_split import IndexFullTextSearchSplit + from pypaimon.table.source.search_types import FullTextScanPlan + from pypaimon.tests.vector_search_filter_test import _StubTable, _entry, _field, match_query + + fields = [_field(0, 'title', 'STRING'), _field(1, 'body', 'STRING')] + entries = [_entry(None, field.id, 'full-text', field.name, 0, 1) for field in fields] + table = _StubTable(fields, entries) + table.options = CoreOptions(Options({'data-evolution.enabled': 'true'})) + files_by_column = {field.name: entry.index_file for field, entry in zip(fields, entries)} + splits = [IndexFullTextSearchSplit(column, 0, 1, [files_by_column[column]]) for column in columns] + query = match_query('needle') + read = DataEvolutionFullTextRead(table, 1, fields[1], query) + executor = InlineExecutor() + + class IndexReader: + def __init__(self, name): + self.name = name + + def visit_full_text_search(self, search): + assert search.field_name == 'body' and search.query == query + return _completed_future(DictBasedScoredIndexResult( + {0: 0.9} if self.name == 'title' else {1: 0.7})) + + def close(self): + pass + + with patch('pypaimon.execution.python.read.search.full_text_read._create_full_text_reader', + side_effect=lambda kind, io, path, metas, bound: IndexReader(metas[0].file_name)): + plan = FullTextScanPlan(splits) + local = read.read_plan(plan) + scheduled = read.read_plan(plan, executor=executor) + assert list(local.results()) == list(scheduled.results()) == expected + assert [task.item.column_name for task in executor.tasks] == (['body'] if expected else []) + + +@pytest.mark.parametrize('family', ['single', 'batch', 'full-text']) +@pytest.mark.parametrize('failure_type', [None, ValueError, CancelledError, KeyboardInterrupt]) +def test_python_scheduled_read_preserves_primary_and_surfaces_cleanup_errors(family, failure_type): + from pypaimon.execution.python.read.search.full_text_read import DataEvolutionFullTextRead + from pypaimon.execution.python.read.search.vector_search_read import ( + BatchVectorSearchReadImpl, DataEvolutionVectorRead, + ) + from pypaimon.table.source.full_text_search_split import IndexFullTextSearchSplit + from pypaimon.table.source.search_types import FullTextScanPlan, VectorSearchScanPlan + from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit + + table = SimpleNamespace(options=SimpleNamespace(data_evolution_enabled=lambda: True)) + table.copy = lambda options: table + column = SimpleNamespace(name='value') + index = SimpleNamespace(index_type='full-text' if family == 'full-text' else 'ivf-flat') + if family == 'full-text': + reader = DataEvolutionFullTextRead(table, 1, column, 'needle') + plan = FullTextScanPlan([IndexFullTextSearchSplit('value', 0, 1, [index])]) + execute = reader.read_plan + result = {1: 0.5} + else: + reader = (BatchVectorSearchReadImpl(table, 1, column, [[1.]]) if family == 'batch' + else DataEvolutionVectorRead(table, 1, column, [1.])) + plan = VectorSearchScanPlan([IndexVectorSearchSplit(0, 1, [index])]) + execute = reader.read_batch_plan if family == 'batch' else reader.read_plan + result = (None, [{1: 0.5}] if family == 'batch' else {1: 0.5}) + + primary = None if failure_type is None else failure_type('worker failed') + cleanup = RuntimeError('scheduler cleanup failed') + + class Results: + def __init__(self): + self.returned = False + self.closed = 0 + + def __iter__(self): + return self + + def __next__(self): + if primary is not None: + raise primary + if self.returned: + raise StopIteration + self.returned = True + return 0, result + + def close(self): + self.closed += 1 + raise cleanup + + results = Results() + with pytest.raises(BaseException) as caught: + execute(plan, executor=MapExecutor(lambda tasks, ordered: results)) + assert caught.value is (cleanup if primary is None else primary) + assert results.closed == 1 From b143a8fb759cb25a506be31b0c8772f8f481f9ad Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 22 Oct 2026 09:01:00 +0800 Subject: [PATCH 6/9] Fix Python CI checks and preserve writer configuration contracts --- .github/workflows/ci-python.yml | 3 +- paimon-python/conftest.py | 2 + .../execution/native/write/native_write.py | 4 +- .../pypaimon/tests/blob_table_test.py | 38 +++++++++++---- .../tests/data_file_path_directory_test.py | 17 ++++--- .../tests/datasource_ray_adapter_test.py | 2 +- .../tests/lazy_search_execution_test.py | 2 +- .../tests/native_test_instrumentation_test.py | 30 ++++++++++++ .../pypaimon/tests/native_write_test.py | 4 +- .../table_update_arrow_compatibility_test.py | 7 +-- paimon-python/pypaimon/tests/variant_test.py | 2 +- .../tests/write/backend_boundary_test.py | 46 +++++++++++++++++++ .../pypaimon/write/write_dispatch.py | 12 +++-- 13 files changed, 138 insertions(+), 31 deletions(-) diff --git a/.github/workflows/ci-python.yml b/.github/workflows/ci-python.yml index b14c3d5424dc..775c578943fb 100644 --- a/.github/workflows/ci-python.yml +++ b/.github/workflows/ci-python.yml @@ -312,7 +312,8 @@ jobs: compatibility: name: Package and Ray compatibility timeout-minutes: 60 - runs-on: ubuntu-latest + # setup-python provides Python 3.9 builds for Ubuntu 24.04, not 26.04. + runs-on: ubuntu-24.04 steps: - name: Checkout code uses: actions/checkout@v6 diff --git a/paimon-python/conftest.py b/paimon-python/conftest.py index b36b91d6f09d..394fcda7e345 100644 --- a/paimon-python/conftest.py +++ b/paimon-python/conftest.py @@ -15,6 +15,7 @@ # limitations under the License. import os +from functools import wraps import pytest @@ -139,6 +140,7 @@ def tracked_read(self, splits, *args, **kwargs): def track_write(method): original_write = getattr(NativeTableWrite, method) + @wraps(original_write) def tracked_write(self, data, *args, **kwargs): global _native_write_count native = self._native_writer is not None diff --git a/paimon-python/pypaimon/execution/native/write/native_write.py b/paimon-python/pypaimon/execution/native/write/native_write.py index 18a08b0d038d..7d726c3c7cd1 100644 --- a/paimon-python/pypaimon/execution/native/write/native_write.py +++ b/paimon-python/pypaimon/execution/native/write/native_write.py @@ -130,9 +130,9 @@ def consume(field_name, encoded): return self def with_blob_uri_reader_factory(self, uri_reader_factory): + if uri_reader_factory is not None and not callable(getattr(uri_reader_factory, 'create', None)): + raise TypeError('URI reader factory must provide callable create') if self._blob_rows is not None: - if uri_reader_factory is not None and not callable(getattr(uri_reader_factory, 'create', None)): - raise TypeError('URI reader factory must provide callable create') # Only core knows whether RowKind filtering accepted any rows or a # postpone writer buffered data. Validate before changing the live # bridge's fallback, including when clearing the application factory. diff --git a/paimon-python/pypaimon/tests/blob_table_test.py b/paimon-python/pypaimon/tests/blob_table_test.py index daa02bee19e1..a972c040de45 100755 --- a/paimon-python/pypaimon/tests/blob_table_test.py +++ b/paimon-python/pypaimon/tests/blob_table_test.py @@ -6291,13 +6291,17 @@ def my_consumer(field_name, descriptor): data = Blob.from_descriptor(uri_reader, desc).to_data() self.assertEqual(data, blob_bytes) - def test_blob_consumer_after_write_raises(self): - """Setting consumer after data has been written must raise.""" + def test_blob_consumer_changes_apply_to_new_partition_writers(self): + """Existing writers retain their consumer when the factory changes.""" + from pypaimon.common.uri_reader import FileUriReader + from pypaimon.table.row.blob import Blob + pa_schema = pa.schema([ ('id', pa.int32()), + ('pt', pa.int32()), ('blob_data', pa.large_binary()), ]) - schema = Schema.from_pyarrow_schema(pa_schema, options={ + schema = Schema.from_pyarrow_schema(pa_schema, partition_keys=['pt'], options={ 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true', }) @@ -6306,15 +6310,29 @@ def test_blob_consumer_after_write_raises(self): write_builder = table.new_batch_write_builder() writer = write_builder.new_write() + commit = write_builder.new_commit() + received = [] - writer.write_arrow(pa.Table.from_pydict({ - 'id': [1], - 'blob_data': [b'data'], - }, schema=pa_schema)) + def consumer(field, descriptor): + received.append((field, descriptor)) + return False - with self.assertRaises(RuntimeError): - writer.with_blob_consumer(lambda f, d: False) - writer.close() + try: + writer.write_arrow(pa.Table.from_pydict({ + 'id': [1], 'pt': [1], 'blob_data': [b'old'], + }, schema=pa_schema)) + self.assertIs(writer.with_blob_consumer(consumer), writer) + writer.write_arrow(pa.Table.from_pydict({ + 'id': [2, 3], 'pt': [1, 2], 'blob_data': [b'existing', b'new'], + }, schema=pa_schema)) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + self.assertEqual(len(received), 1) + self.assertEqual(received[0][0], 'blob_data') + blob = Blob.from_descriptor(FileUriReader(table.file_io), received[0][1]) + self.assertEqual(blob.to_data(), b'new') if __name__ == '__main__': diff --git a/paimon-python/pypaimon/tests/data_file_path_directory_test.py b/paimon-python/pypaimon/tests/data_file_path_directory_test.py index 80addf9a815e..b6e9e4fe5f9f 100644 --- a/paimon-python/pypaimon/tests/data_file_path_directory_test.py +++ b/paimon-python/pypaimon/tests/data_file_path_directory_test.py @@ -111,7 +111,9 @@ def test_native_write_dispatches_with_configured_directory(self): "write.native.enabled": "true"}) from pypaimon.write.write_dispatch import NativeTableWrite - adapter = mock.Mock(spec=['write_arrow_batch', 'prepare_commit', 'close']) + adapter = mock.Mock(spec=['write_arrow_batch', 'prepare_commit', 'close', 'stream', '_batch_committed']) + adapter.stream = False + adapter._batch_committed = False messages = [] adapter.prepare_commit.return_value = messages builder = table.new_batch_write_builder() @@ -129,7 +131,8 @@ def test_native_write_dispatches_with_configured_directory(self): writer.write_arrow_batch(batch) adapter.write_arrow_batch.assert_called_once_with(batch, None) self.assertIs(writer.prepare_commit(), messages) - adapter.prepare_commit.assert_called_once_with(None) + from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER + adapter.prepare_commit.assert_called_once_with(BATCH_COMMIT_IDENTIFIER) finally: writer.close() adapter.close.assert_called_once_with() @@ -139,7 +142,7 @@ def test_native_commit_dispatches_with_configured_directory(self): "db.t_native_commit", {"data-file.path-directory": "data", "commit.native.enabled": "true"}) commit = table.new_batch_write_builder().new_commit() - sentinel = object() + sentinel = mock.Mock() try: with mock.patch( "pypaimon.execution.native.write.native_commit.native_messages_supported", return_value=True), \ @@ -147,11 +150,13 @@ def test_native_commit_dispatches_with_configured_directory(self): "pypaimon.execution.native.write.native_commit.create_native_commit", return_value=sentinel), \ mock.patch( "pypaimon.execution.native.write.native_commit.to_native_commit_messages", return_value=[]): - self.assertEqual(commit._prepare_native_commit([]), (sentinel, [])) + prepared, messages = commit._prepare_native_commit([]) + self.assertIs(prepared._native_commit, sentinel) + self.assertFalse(prepared._is_batch) + self.assertEqual(messages, []) finally: - # The sentinel represents a prepared committer, not a real resource. - commit._native_commit = None commit.close() + sentinel.close.assert_called_once_with() def test_default_keeps_data_files_at_bucket_root(self): self.catalog.create_table( diff --git a/paimon-python/pypaimon/tests/datasource_ray_adapter_test.py b/paimon-python/pypaimon/tests/datasource_ray_adapter_test.py index 1cdcc36d2a1f..7dd6c11b8cb3 100644 --- a/paimon-python/pypaimon/tests/datasource_ray_adapter_test.py +++ b/paimon-python/pypaimon/tests/datasource_ray_adapter_test.py @@ -171,7 +171,7 @@ def _distribute_splits_into_equal_chunks(self, splits, n_chunks): assert len(tasks) == 1 assert tasks[0].metadata.num_rows == 2 assert tasks[0].metadata.size_bytes == 10 - assert tasks[0].metadata.input_files == ['file.parquet'] + assert list(tasks[0].metadata.input_files) == ['file.parquet'] def test_public_datasource_configuration_is_lazy_and_pickles_at_original_paths(): diff --git a/paimon-python/pypaimon/tests/lazy_search_execution_test.py b/paimon-python/pypaimon/tests/lazy_search_execution_test.py index 8b91fbb14426..36313fd9e671 100644 --- a/paimon-python/pypaimon/tests/lazy_search_execution_test.py +++ b/paimon-python/pypaimon/tests/lazy_search_execution_test.py @@ -461,7 +461,7 @@ def test_python_source_metadata_round_trips_through_public_native_plan_construct assert isinstance(native.splits()[0], datafusion.IndexFullTextSearchSplit if full_text else datafusion.IndexVectorSearchSplit) assert copied.snapshot()._to_json() == native.snapshot()._to_json() - assert type(copied.splits()[0]) is type(native.splits()[0]) + assert copied.splits()[0].__class__ is native.splits()[0].__class__ assert restored.snapshot() == snapshot restored_index = restored.splits()[0] restored_files = restored_index.full_text_index_files if full_text else restored_index.vector_index_files diff --git a/paimon-python/pypaimon/tests/native_test_instrumentation_test.py b/paimon-python/pypaimon/tests/native_test_instrumentation_test.py index 75bd44b8f921..91c5153a111d 100644 --- a/paimon-python/pypaimon/tests/native_test_instrumentation_test.py +++ b/paimon-python/pypaimon/tests/native_test_instrumentation_test.py @@ -18,7 +18,9 @@ """Native CI instrumentation preserves the public mutation API.""" import importlib.util +import inspect from pathlib import Path +from types import SimpleNamespace from unittest.mock import Mock import pytest @@ -30,6 +32,34 @@ pytest.mark.python_write, pytest.mark.python_commit] +def test_native_write_tracking_preserves_signature_and_forwards_bucket(monkeypatch): + from pypaimon.execution.native.write.native_write import NativeTableWrite + + path = Path(__file__).resolve().parents[2] / 'conftest.py' + spec = importlib.util.spec_from_file_location('native_write_instrumentation', str(path)) + instrumentation = importlib.util.module_from_spec(spec) + spec.loader.exec_module(instrumentation) + for key in ('PLAN', 'READ', 'COMMIT', 'UPDATE'): + monkeypatch.setenv('PYPAIMON_TEST_NATIVE_' + key, '0') + monkeypatch.setenv('PYPAIMON_TEST_NATIVE_WRITE', '1') + calls = [] + result = object() + + def write_arrow_batch(self, data, bucket=None): + calls.append((data, bucket)) + return result + + monkeypatch.setattr(NativeTableWrite, 'write_arrow_batch', write_arrow_batch) + instrumentation.pytest_configure(Mock()) + instrumentation._force_native_write_for_test = True + assert inspect.signature(NativeTableWrite.write_arrow_batch) == inspect.signature(write_arrow_batch) + data = SimpleNamespace(num_rows=1) + writer = SimpleNamespace(_native_writer=object()) + assert NativeTableWrite.write_arrow_batch(writer, data, bucket=7) is result + assert calls == [(data, 7)] + assert instrumentation._native_write_count == 1 + + @pytest.mark.parametrize('operation,args', [ ('truncate_table', ()), ('truncate_partitions', ([{'pt': 'a'}],)), ]) diff --git a/paimon-python/pypaimon/tests/native_write_test.py b/paimon-python/pypaimon/tests/native_write_test.py index 625cb73f0fd0..9f0e241d0a7b 100644 --- a/paimon-python/pypaimon/tests/native_write_test.py +++ b/paimon-python/pypaimon/tests/native_write_test.py @@ -334,9 +334,9 @@ def test_writer_abort_preserves_native_commit_files(native_rest_catalog, failure native, native_messages = prepared proxy = Mock(wraps=native) - def publish(*args): + def publish(*args, **kwargs): if failure != 'before-publication': - native.commit(*args) + native.commit(*args, **kwargs) if failure: raise RuntimeError('Native commit outcome is unknown') diff --git a/paimon-python/pypaimon/tests/table_update_arrow_compatibility_test.py b/paimon-python/pypaimon/tests/table_update_arrow_compatibility_test.py index 8e24201cc006..28e72a054ada 100644 --- a/paimon-python/pypaimon/tests/table_update_arrow_compatibility_test.py +++ b/paimon-python/pypaimon/tests/table_update_arrow_compatibility_test.py @@ -63,10 +63,11 @@ def test_predicate_assignment_struct_cast(wrap_type, wrap_value, chunked, projec else: try: expected = values.cast(target).to_pylist() - except pa.ArrowNotImplementedError: - # Older Arrow versions do not implement STRUCT casts. - with pytest.raises(pa.ArrowNotImplementedError): + except (pa.ArrowNotImplementedError, pa.ArrowTypeError) as error: + # Older Arrow versions reject unsupported STRUCT field projections. + with pytest.raises(type(error)) as raised: BatchTableUpdate._assignment_to_array(values, target, 2) + assert str(raised.value) == str(error) return result = BatchTableUpdate._assignment_to_array(values, target, 2) assert result.type == target diff --git a/paimon-python/pypaimon/tests/variant_test.py b/paimon-python/pypaimon/tests/variant_test.py index 46903024482b..80c3bf905987 100644 --- a/paimon-python/pypaimon/tests/variant_test.py +++ b/paimon-python/pypaimon/tests/variant_test.py @@ -1516,7 +1516,7 @@ def typed(value): builder = table.new_batch_write_builder() writer = builder.new_write() self.assertEqual(type(writer).__name__, - 'NativeTableWrite' if native else 'BatchTableWrite') + 'NativeBatchTableWrite' if native else 'BatchTableWrite') commit = builder.new_commit() try: for start in range(0, len(gvs), 3): diff --git a/paimon-python/pypaimon/tests/write/backend_boundary_test.py b/paimon-python/pypaimon/tests/write/backend_boundary_test.py index 452f84e3c632..7af31c2cf735 100644 --- a/paimon-python/pypaimon/tests/write/backend_boundary_test.py +++ b/paimon-python/pypaimon/tests/write/backend_boundary_test.py @@ -110,6 +110,52 @@ def test_invalid_blob_configuration_keeps_the_configured_consumer(tmp_path): writer.close() +@pytest.mark.parametrize('row_bridge', [False, True]) +def test_native_uri_reader_factory_validation_precedes_core_configuration(tmp_path, row_bridge): + core = Mock() + writer = NativeWrite(_table(tmp_path), 'job', None, False, core) + factory = Mock() + if row_bridge: + writer._blob_rows = Mock() + try: + assert writer.with_blob_uri_reader_factory(factory) is writer + core.with_blob_uri_reader_factory.reset_mock() + with pytest.raises(TypeError, match='factory must provide callable create'): + writer.with_blob_uri_reader_factory(42) + assert writer._blob_uri_reader_factory is factory + if row_bridge: + assert writer._blob_rows.fallback is factory + core.with_blob_uri_reader_factory.assert_not_called() + finally: + writer.close() + + +@pytest.mark.parametrize('method,attribute', [ + ('with_write_type', 'write_cols'), + ('with_blob_consumer', 'blob_consumer'), + ('with_blob_uri_reader_factory', 'blob_uri_reader_factory'), +]) +def test_native_facade_configuration_keeps_identity_after_python_selection(tmp_path, method, attribute): + from pypaimon.write.write_dispatch import NativeTableWrite + + table = _table(tmp_path) + core = Mock() + writer = NativeTableWrite(table, 'job', None, False, core) + try: + store = writer.file_store_write + value = ['value'] if method == 'with_write_type' else Mock() + reset = table.field_names if method == 'with_write_type' else None + configure = getattr(writer, method) + assert configure(value) is writer + assert getattr(store, attribute) == value + assert configure(reset) is writer + assert getattr(store, attribute) is None + getattr(core, method).assert_not_called() + finally: + writer.close() + core.close.assert_called_once_with() + + def test_python_delete_matching_respects_independent_read_backend(tmp_path): table = _table(tmp_path, {'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true', 'deletion-vectors.enabled': 'true'}) diff --git a/paimon-python/pypaimon/write/write_dispatch.py b/paimon-python/pypaimon/write/write_dispatch.py index e7c80dacfbad..fb0d2910caf3 100644 --- a/paimon-python/pypaimon/write/write_dispatch.py +++ b/paimon-python/pypaimon/write/write_dispatch.py @@ -149,23 +149,27 @@ def batch_committed(self, value): def with_write_type(self, write_cols): names = validate_write_columns(write_cols) if self._python_writer is not None: - return self._python_writer.with_write_type(names) + self._python_writer.with_write_type(names) + return self full_names = self._table_schema.names if (self.table.is_primary_key_table and not (len(names) == len(full_names) and set(names) == set(full_names))): - return self._switch_to_python().with_write_type(names) + self._switch_to_python().with_write_type(names) + return self self._execution.with_write_type(names) return self def with_blob_consumer(self, blob_consumer): if self._python_writer is not None: - return self._python_writer.with_blob_consumer(blob_consumer) + self._python_writer.with_blob_consumer(blob_consumer) + return self self._execution.with_blob_consumer(blob_consumer) return self def with_blob_uri_reader_factory(self, uri_reader_factory): if self._python_writer is not None: - return self._python_writer.with_blob_uri_reader_factory(uri_reader_factory) + self._python_writer.with_blob_uri_reader_factory(uri_reader_factory) + return self self._execution.with_blob_uri_reader_factory(uri_reader_factory) return self From 91611e6e8af2a8e2cc6c95ef6889ffad01d00eda Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 22 Oct 2026 09:44:50 +0800 Subject: [PATCH 7/9] [python] Align idle append writer retirement with Java --- .../python/write/file_store_write.py | 85 ++++ .../execution/python/write/table_write.py | 2 +- .../tests/write/append_row_kind_test.py | 2 + .../write/append_writer_lifecycle_test.py | 380 ++++++++++++++++++ 4 files changed, 468 insertions(+), 1 deletion(-) create mode 100644 paimon-python/pypaimon/tests/write/append_writer_lifecycle_test.py diff --git a/paimon-python/pypaimon/execution/python/write/file_store_write.py b/paimon-python/pypaimon/execution/python/write/file_store_write.py index 497dd50a9531..4e687254d7a3 100644 --- a/paimon-python/pypaimon/execution/python/write/file_store_write.py +++ b/paimon-python/pypaimon/execution/python/write/file_store_write.py @@ -36,6 +36,7 @@ logger = logging.getLogger(__name__) +_MIN_COMMIT_IDENTIFIER = -(1 << 63) class FileStoreWrite: @@ -55,12 +56,18 @@ def __init__(self, table, commit_user): check_supported(table) self.table: FileStoreTable = table + self._commit_user = commit_user self.data_writers: Dict[Tuple, DataWriter] = {} # Like Java's per-bucket writer factory, existing writers retain the # requested layout even when the factory changes for future buckets. # DataWriter.write_cols may instead describe optimized physical files. self._write_columns_by_bucket = {} self._runtime_total_buckets: Dict[Tuple, int] = {} + self._last_modified_commit_identifiers = {} + self._written_buckets = set() + # Factory lifecycle guards survive idle retirement of every bucket. + self._ever_written = False + self._prepared_writer_lifecycle = None self.max_seq_numbers: dict = {} self.restore_snapshot_id = None self.write_cols = None @@ -102,6 +109,7 @@ def write( if write_columns is not None: data = project_arrow_batch(data, write_columns) writer.write(data) + self._written_buckets.add(key) def _new_bucket_writer(self, partition, bucket): columns = tuple(self.write_cols) if self.write_cols is not None else None @@ -109,6 +117,7 @@ def _new_bucket_writer(self, partition, bucket): key = (partition, bucket) self.data_writers[key] = writer self._write_columns_by_bucket[key] = columns + self._ever_written = True def _write_columns_for(self, partition, bucket): return self._write_columns_by_bucket.get((partition, bucket), self.write_cols) @@ -153,6 +162,7 @@ def write_row( writer = self.data_writers[key] if hasattr(writer, 'write_row'): writer.write_row(row) + self._written_buckets.add(key) return data = row_values_to_arrow_table( @@ -163,6 +173,7 @@ def write_row( from pypaimon.write.row_kind import with_row_kind data = with_row_kind(self.table, data, row) writer.write(data.to_batches()[0]) + self._written_buckets.add(key) def roll_before_group_if_needed(self, row_count: int): for writer in self.data_writers.values(): @@ -374,11 +385,17 @@ def prepare_commit(self, commit_identifier) -> List[CommitMessage]: def _prepare_commit_messages(self, commit_identifier) -> List[CommitMessage]: """Collect files while their parent is still preparing the complete increment.""" self.commit_identifier = commit_identifier + self._prepared_writer_lifecycle = None + # Resolve metadata before draining any writer, so lookup failures leave + # every partition's pending rows and prepared files available for retry. + idle_writers = self._idle_append_writers() commit_messages = [] + modified_writers = set() for (partition, bucket), writer in self.data_writers.items(): committed_files = writer.prepare_commit() changelog_files = writer.prepare_changelog_commit() if committed_files or changelog_files: + modified_writers.add((partition, bucket)) commit_message = CommitMessage( partition=partition, bucket=bucket, @@ -387,13 +404,75 @@ def _prepare_commit_messages(self, commit_identifier) -> List[CommitMessage]: total_buckets=self._runtime_total_buckets.get(partition), ) commit_messages.append(commit_message) + self._prepared_writer_lifecycle = ( + commit_identifier, modified_writers, idle_writers - modified_writers) return commit_messages + def _idle_append_writers(self): + if self.table.is_primary_key_table or not any( + identifier != _MIN_COMMIT_IDENTIFIER + for identifier in self._last_modified_commit_identifiers.values()): + # Like Java, the first prepare must retain all newly opened writers. + return set() + mode = self.table.bucket_mode() + if mode not in (BucketMode.BUCKET_UNAWARE, BucketMode.HASH_FIXED): + return set() + idle = set(self.data_writers).difference(self._written_buckets) + if not idle or mode == BucketMode.BUCKET_UNAWARE: + return idle + latest_identifier = self._latest_committed_identifier_from_filesystem() + if latest_identifier is None: + return set() + return {key for key in idle + if self._last_modified_commit_identifiers.get(key, _MIN_COMMIT_IDENTIFIER) < latest_identifier} + + def _latest_committed_identifier_from_filesystem(self): + # Java FileSystemWriteRestore uses filesystem snapshots for writer + # cleanup, independently of catalog loaders and a caller's read view. + manager = self.table.snapshot_manager() + try: + snapshot = manager._get_latest_snapshot_from_filesystem() + except FileNotFoundError: + return None + while snapshot is not None: + if snapshot.commit_user == self._commit_user: + return snapshot.commit_identifier + if snapshot.id <= 1: + break + try: + snapshot = manager.get_snapshot_by_id(snapshot.id - 1) + except FileNotFoundError: + # Expiration can race the exists/read pair. A gap ends the + # search; permission, transport and decoding errors propagate. + break + return None + def _release_prepared_files(self): + lifecycle = self._prepared_writer_lifecycle + if lifecycle is not None: + _, _, retired = lifecycle + # All data/index preparations succeeded. Close idle writers before + # releasing metadata, so a close failure can still retry the files + # prepared by active writers in this checkpoint. + for key in retired: + self.data_writers[key].close() # Hand off only after every partition prepared successfully. Until # then, close/abort must still clean up files from earlier partitions. for writer in self.data_writers.values(): writer._release_prepared_files() + if lifecycle is not None: + identifier, modified, retired = lifecycle + for key in modified: + self._last_modified_commit_identifiers[key] = identifier + for key in retired: + del self.data_writers[key] + self._write_columns_by_bucket.pop(key, None) + self._last_modified_commit_identifiers.pop(key, None) + partitions = {partition for partition, _ in self.data_writers} + for partition in set(self._runtime_total_buckets).difference(partitions): + del self._runtime_total_buckets[partition] + self._written_buckets.clear() + self._prepared_writer_lifecycle = None # Existing bucket writers keep their advancing sequence counters. A # bucket opened after this checkpoint must restore from the latest # snapshot rather than reuse another bucket's earlier partition scan. @@ -406,6 +485,9 @@ def close(self): self.data_writers.clear() self._write_columns_by_bucket.clear() self._runtime_total_buckets.clear() + self._last_modified_commit_identifiers.clear() + self._written_buckets.clear() + self._prepared_writer_lifecycle = None def abort(self): """Abort all data writers and clean up files produced by this write.""" @@ -417,6 +499,9 @@ def abort(self): self.data_writers.clear() self._write_columns_by_bucket.clear() self._runtime_total_buckets.clear() + self._last_modified_commit_identifiers.clear() + self._written_buckets.clear() + self._prepared_writer_lifecycle = None def _seq_number_stats(self, partition: Tuple) -> Dict[int, int]: buckets = self.max_seq_numbers.get(partition) diff --git a/paimon-python/pypaimon/execution/python/write/table_write.py b/paimon-python/pypaimon/execution/python/write/table_write.py index 505ef95748cd..c3d4e1016557 100644 --- a/paimon-python/pypaimon/execution/python/write/table_write.py +++ b/paimon-python/pypaimon/execution/python/write/table_write.py @@ -234,7 +234,7 @@ def with_blob_uri_reader_factory(self, uri_reader_factory): writer = self.writer if uri_reader_factory is not None and not callable(getattr(uri_reader_factory, 'create', None)): raise TypeError('URI reader factory must provide callable create') - if writer.file_store_write.data_writers: + if writer.file_store_write._ever_written: raise RuntimeError( "with_blob_uri_reader_factory must be called before any " "write operation." diff --git a/paimon-python/pypaimon/tests/write/append_row_kind_test.py b/paimon-python/pypaimon/tests/write/append_row_kind_test.py index e97c4946407f..b6204e13d409 100644 --- a/paimon-python/pypaimon/tests/write/append_row_kind_test.py +++ b/paimon-python/pypaimon/tests/write/append_row_kind_test.py @@ -365,6 +365,8 @@ def column(self, index): writer.data_writers = {} writer._write_columns_by_bucket = {} writer._runtime_total_buckets = {} + writer._written_buckets = set() + writer._ever_written = False writer.options = object() target = Mock() writer._create_data_writer = Mock(return_value=target) diff --git a/paimon-python/pypaimon/tests/write/append_writer_lifecycle_test.py b/paimon-python/pypaimon/tests/write/append_writer_lifecycle_test.py new file mode 100644 index 000000000000..ea0bbceb3816 --- /dev/null +++ b/paimon-python/pypaimon/tests/write/append_writer_lifecycle_test.py @@ -0,0 +1,380 @@ +# 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. + +"""Append writers capture factory settings until Java's idle retirement boundary.""" + +import gc +import weakref +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pyarrow as pa +import pytest + +from pypaimon import CatalogFactory, Schema +from pypaimon.schema.data_types import AtomicType, DataField, PyarrowFieldParser +from pypaimon.tests.native_blob_uri_reader_test import Factory +from pypaimon.tests.write.write_factory_reconfiguration_test import _ENGINES, _read, _table, _write + +pytestmark = [pytest.mark.python_plan, pytest.mark.python_read, + pytest.mark.python_write, pytest.mark.python_commit] + + +@pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) +@pytest.mark.parametrize('idle_prepares', [0, 1, 2]) +@pytest.mark.parametrize('mode', ['arrow', 'row']) +def test_unaware_idle_prepare_reopens_with_current_write_type(tmp_path, native, idle_prepares, mode): + table = _table(tmp_path, native) + builder = table.new_stream_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + assert writer.prepare_commit(0) == [] + writer.with_write_type(['id', 'pt']) + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}], mode) + commit.commit(writer.prepare_commit(1), 1) + for identifier in range(2, 2 + idle_prepares): + assert writer.prepare_commit(identifier) == [] + writer.with_write_type(['value', 'pt']) + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}], mode) + identifier = 2 + idle_prepares + messages = writer.prepare_commit(identifier) + assert [file.write_cols for message in messages for file in message.new_files] == [ + ['value', 'pt'] if idle_prepares else ['id', 'pt']] + commit.commit(messages, identifier) + finally: + writer.close() + commit.close() + expected = [ + {'id': None, 'value': 20, 'pt': 'old'}, {'id': 1, 'value': None, 'pt': 'old'} + ] if idle_prepares else [ + {'id': 1, 'value': None, 'pt': 'old'}, {'id': 2, 'value': None, 'pt': 'old'}] + assert _read(table) == expected + + +@pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) +@pytest.mark.parametrize('first_identifier', [-(1 << 63), 0], ids=['minimum', 'ordinary']) +def test_minimum_identifier_remains_first_prepare_until_a_later_nonempty_increment( + tmp_path, native, first_identifier): + table = _table(tmp_path, native) + builder = table.new_stream_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.with_write_type(['id', 'pt']) + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}]) + commit.commit(writer.prepare_commit(first_identifier), first_identifier) + assert writer.prepare_commit(first_identifier + 1) == [] + writer.with_write_type(['value', 'pt']) + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}]) + messages = writer.prepare_commit(first_identifier + 1) + minimum = first_identifier == -(1 << 63) + assert [file.write_cols for message in messages for file in message.new_files] == [ + ['id', 'pt'] if minimum else ['value', 'pt']] + commit.commit(messages, first_identifier + 1) + # Once a nonempty increment has advanced past MIN, this idle bucket + # can retire and capture the factory's new full layout. + assert writer.prepare_commit(first_identifier + 2) == [] + writer.with_write_type(table.field_names) + _write(writer, table, [{'id': 3, 'value': 30, 'pt': 'old'}]) + messages = writer.prepare_commit(first_identifier + 3) + assert [file.write_cols for message in messages for file in message.new_files] == [None] + commit.commit(messages, first_identifier + 3) + finally: + writer.close() + commit.close() + expected = [{'id': 1, 'value': None, 'pt': 'old'}, {'id': 3, 'value': 30, 'pt': 'old'}] + expected.append({'id': 2, 'value': None, 'pt': 'old'} if minimum + else {'id': None, 'value': 20, 'pt': 'old'}) + assert _read(table) == sorted(expected, key=lambda row: row['id'] or -1) + + +@pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) +@pytest.mark.parametrize('factory', [None, Factory()], ids=['clear', 'set']) +def test_idle_retirement_does_not_reopen_uri_reader_factory_configuration(tmp_path, native, factory): + table = _table(tmp_path, native) + builder = table.new_stream_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.with_blob_uri_reader_factory(Factory()) + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}]) + commit.commit(writer.prepare_commit(1), 1) + assert writer.prepare_commit(2) == [] + if not native: + assert writer.file_store_write.data_writers == {} + with pytest.raises(RuntimeError, match='before any write'): + writer.with_blob_uri_reader_factory(factory) + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}]) + commit.commit(writer.prepare_commit(3), 3) + finally: + writer.close() + commit.close() + assert _read(table) == [ + {'id': 1, 'value': 10, 'pt': 'old'}, {'id': 2, 'value': 20, 'pt': 'old'}] + + +@pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) +@pytest.mark.parametrize('input_', ['row', 'batch', 'table', 'pandas']) +def test_ignored_events_do_not_lock_uri_reader_factory_configuration(tmp_path, native, input_): + from pypaimon.tests.write.append_row_kind_test import _table as event_table, _write as write_event + table, schema = event_table(tmp_path, native, **{'ignore-delete': 'true'}) + writer = table.new_stream_write_builder().new_write() + try: + write_event(writer, table, schema, [dict(id=1, op='-D', pt='old')], input_, True) + assert writer.prepare_commit(1) == [] + assert writer.with_blob_uri_reader_factory(Factory()) is writer + assert writer.with_blob_uri_reader_factory(None) is writer + write_event(writer, table, schema, [dict(id=2, op='+I', pt='old')], input_, True) + with pytest.raises(RuntimeError, match='before any write'): + writer.with_blob_uri_reader_factory(None) + finally: + writer.close() + + +@pytest.mark.parametrize('failure_stage', ['create', 'write']) +def test_uri_reader_factory_locks_once_bucket_creation_succeeds(tmp_path, failure_stage): + table = _table(tmp_path, False) + writer = table.new_stream_write_builder().new_write() + failure = OSError('bucket %s failed' % failure_stage) + creation = ({'side_effect': failure} if failure_stage == 'create' + else {'return_value': Mock(write=Mock(side_effect=failure))}) + try: + with patch.object(writer.file_store_write, '_create_data_writer', **creation): + with pytest.raises(OSError) as caught: + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}]) + assert caught.value is failure + if failure_stage == 'create': + assert writer.with_blob_uri_reader_factory(Factory()) is writer + else: + with pytest.raises(RuntimeError, match='before any write'): + writer.with_blob_uri_reader_factory(Factory()) + finally: + writer.close() + + +@pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) +@pytest.mark.parametrize('latest_commit', ['same-id', 'other-user', 'higher-id']) +def test_fixed_idle_prepare_requires_strictly_newer_commit_of_same_user(tmp_path, native, latest_commit): + table = _table(tmp_path, native, precomputed=True) + builder = table.new_stream_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.with_write_type(['id', 'pt']) + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}], 'precomputed') + commit.commit(writer.prepare_commit(1), 1) + # Equality is insufficient, even if another checkpoint is being prepared. + assert writer.prepare_commit(2) == [] + expected = [{'id': 1, 'value': None, 'pt': 'old'}] + if latest_commit == 'other-user': + other = table.new_stream_write_builder() + other_writer, other_commit = other.new_write(), other.new_commit() + try: + _write(other_writer, table, [{'id': 8, 'value': 80, 'pt': 'helper'}], 'precomputed') + other_commit.commit(other_writer.prepare_commit(99), 99) + finally: + other_writer.close() + other_commit.close() + expected.insert(0, {'id': 8, 'value': 80, 'pt': 'helper'}) + elif latest_commit == 'higher-id': + _write(writer, table, [{'id': 8, 'value': 80, 'pt': 'helper'}], 'precomputed') + commit.commit(writer.prepare_commit(2), 2) + expected.insert(0, {'id': 8, 'value': None, 'pt': 'helper'}) + assert writer.prepare_commit(3) == [] + writer.with_write_type(['value', 'pt']) + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}], 'precomputed') + messages = writer.prepare_commit(4) + retired = latest_commit == 'higher-id' + assert [file.write_cols for message in messages for file in message.new_files] == [ + ['value', 'pt'] if retired else ['id', 'pt']] + commit.commit(messages, 4) + finally: + writer.close() + commit.close() + expected.append({'id': None, 'value': 20, 'pt': 'old'} if retired + else {'id': 2, 'value': None, 'pt': 'old'}) + assert _read(table) == sorted(expected, key=lambda row: (row['pt'], row['id'] or -1)) + + +@pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) +def test_idle_retirement_releases_captured_blob_consumer(tmp_path, native): + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.create_database('default', True) + fields = [DataField(0, 'id', AtomicType('INT')), DataField(1, 'payload', AtomicType('BLOB'))] + options = {'write.native.enabled': str(native).lower(), 'commit.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', 'read.native.enabled': 'false', + 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'} + catalog.create_table('default.blobs', Schema(fields=fields, options=options), False) + table = catalog.get_table('default.blobs') + builder = table.new_stream_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + received = [] + + class Consumer: + def __call__(self, name, descriptor): + received.append((name, descriptor)) + return True + + consumer = Consumer() + reference = weakref.ref(consumer) + try: + writer.with_blob_consumer(consumer) + schema = PyarrowFieldParser.from_paimon_schema(fields) + writer.write_arrow(pa.Table.from_pylist([{'id': 1, 'payload': b'old'}], schema=schema)) + commit.commit(writer.prepare_commit(1), 1) + writer.with_blob_consumer(None) + del consumer + gc.collect() + assert reference() is not None + assert len(received) == 1 + assert writer.prepare_commit(2) == [] + gc.collect() + assert reference() is None + writer.write_arrow(pa.Table.from_pylist([{'id': 2, 'payload': b'new'}], schema=schema)) + commit.commit(writer.prepare_commit(3), 3) + assert len(received) == 1 + finally: + writer.close() + commit.close() + reader = table.new_read_builder() + assert sorted(reader.new_read().to_arrow(reader.new_scan().plan().splits()).to_pylist(), + key=lambda row: row['id']) == [ + {'id': 1, 'payload': b'old'}, {'id': 2, 'payload': b'new'}] + + +@pytest.mark.parametrize('missing', [None, FileNotFoundError('expired during read')]) +@pytest.mark.parametrize('lookup', ['latest', 'previous']) +def test_fixed_snapshot_search_stops_at_expired_gap(tmp_path, missing, lookup): + table = _table(tmp_path, False, precomputed=True) + writer = table.new_stream_write_builder().new_write().with_write_type(['id', 'pt']) + try: + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}], 'precomputed') + writer.prepare_commit(1) + latest = SimpleNamespace(id=3, commit_user='other', commit_identifier=99) + options = {'side_effect': missing} if isinstance(missing, Exception) else {'return_value': None} + with patch('pypaimon.snapshot.snapshot_manager.SnapshotManager._get_latest_snapshot_from_filesystem', + **(options if lookup == 'latest' else {'return_value': latest})), patch( + 'pypaimon.snapshot.snapshot_manager.SnapshotManager.get_snapshot_by_id', + **options) as get, patch( + 'pypaimon.snapshot.snapshot_manager.SnapshotManager.get_latest_snapshot', + side_effect=AssertionError('writer cleanup must not use the catalog snapshot loader')): + assert writer.prepare_commit(2) == [] + assert get.call_args_list == ([] if lookup == 'latest' else [((2,), {})]) + writer.with_write_type(['value', 'pt']) + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}], 'precomputed') + messages = writer.prepare_commit(3) + assert [file.write_cols for message in messages for file in message.new_files] == [['id', 'pt']] + finally: + writer.close() + + +@pytest.mark.parametrize('lookup', ['latest', 'previous']) +def test_snapshot_lookup_failure_preserves_all_pending_files_for_retry(tmp_path, lookup): + table = _table(tmp_path, False, precomputed=True) + writer = table.new_stream_write_builder().new_write() + try: + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}], 'precomputed') + writer.prepare_commit(1) + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'new'}], 'precomputed') + failure = OSError('snapshot permission denied') + latest = SimpleNamespace(id=3, commit_user='other', commit_identifier=99) + before = set(tmp_path.rglob('*.parquet')) + with patch('pypaimon.snapshot.snapshot_manager.SnapshotManager._get_latest_snapshot_from_filesystem', + **({'side_effect': failure} if lookup == 'latest' else {'return_value': latest})), patch( + 'pypaimon.snapshot.snapshot_manager.SnapshotManager.get_snapshot_by_id', side_effect=failure): + with pytest.raises(OSError) as caught: + writer.prepare_commit(2) + assert caught.value is failure + assert set(tmp_path.rglob('*.parquet')) == before + messages = writer.prepare_commit(2) + assert [message.partition for message in messages] == [('new',)] + assert [file.row_count for message in messages for file in message.new_files] == [1] + finally: + writer.close() + + +@pytest.mark.parametrize('mode', ['arrow', 'row', 'precomputed']) +def test_fixed_active_writers_do_not_lookup_retirement_snapshots(tmp_path, mode): + table = _table(tmp_path, False, precomputed=True) + writer = table.new_stream_write_builder().new_write() + try: + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}], mode) + writer.prepare_commit(1) + _write(writer, table, [{'id': 1, 'value': 20, 'pt': 'old'}], mode) + with patch('pypaimon.snapshot.snapshot_manager.SnapshotManager._get_latest_snapshot_from_filesystem', + side_effect=AssertionError('active writers need no cleanup lookup')): + messages = writer.prepare_commit(2) + assert [file.row_count for message in messages for file in message.new_files] == [1] + finally: + writer.close() + + +@pytest.mark.parametrize('finish', ['retry', 'abort']) +def test_failed_prepare_does_not_retire_idle_writers_or_release_other_files(tmp_path, finish): + table = _table(tmp_path, False) + writer = table.new_stream_write_builder().new_write().with_write_type(['id', 'pt']) + try: + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}]) + writer.prepare_commit(1) + delivered_files = set(tmp_path.rglob('*.parquet')) + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'new'}, + {'id': 3, 'value': 30, 'pt': 'last'}]) + store = writer.file_store_write + old = store.data_writers[(('old',), 0)] + last = store.data_writers[(('last',), 0)] + failure = OSError('last partition cannot prepare') + with patch.object(last, 'prepare_commit', side_effect=failure), patch.object(old, 'close') as close: + with pytest.raises(OSError) as caught: + writer.prepare_commit(2) + assert caught.value is failure + close.assert_not_called() + assert set(store.data_writers) == {(('old',), 0), (('new',), 0), (('last',), 0)} + if finish == 'abort': + writer.abort() + assert set(tmp_path.rglob('*.parquet')) == delivered_files + return + messages = writer.prepare_commit(2) + assert {message.partition for message in messages} == {('new',), ('last',)} + files = list(tmp_path.rglob('*.parquet')) + assert len(files) == 3 + writer.abort() + assert all(path.exists() for path in files) + finally: + writer.close() + + +def test_retirement_close_failure_keeps_active_files_for_retry(tmp_path): + table = _table(tmp_path, False) + writer = table.new_stream_write_builder().new_write() + try: + _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}]) + writer.prepare_commit(1) + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'new'}]) + store = writer.file_store_write + old = store.data_writers[(('old',), 0)] + failure = OSError('cannot close idle writer') + with patch.object(old, 'close', side_effect=failure): + with pytest.raises(OSError) as caught: + writer.prepare_commit(2) + assert caught.value is failure + messages = writer.prepare_commit(2) + assert [message.partition for message in messages] == [('new',)] + assert [file.row_count for message in messages for file in message.new_files] == [1] + assert set(store.data_writers) == set(store._write_columns_by_bucket) == {(('new',), 0)} + files = list(tmp_path.rglob('*.parquet')) + assert len(files) == 2 + writer.abort() + assert all(path.exists() for path in files) + finally: + writer.close() From ff21edf6356eedfb78dd9f3d992e2edc8a6e1c24 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 22 Oct 2026 10:03:33 +0800 Subject: [PATCH 8/9] [python] Keep one write type and remove per-row schema queries --- docs/docs/pypaimon/writing.md | 17 +- paimon-python/pypaimon/execution/README.md | 19 +- .../execution/native/write/native_write.py | 46 +-- .../python/write/file_store_write.py | 68 ++--- .../execution/python/write/table_write.py | 15 +- ...tive_filtered_writer_configuration_test.py | 7 +- .../tests/write/append_row_kind_test.py | 1 - .../write/append_writer_lifecycle_test.py | 104 ++++--- .../write_factory_reconfiguration_test.py | 262 +++++++++++++----- .../write/writer_execution_protocol_test.py | 94 +++---- 10 files changed, 346 insertions(+), 287 deletions(-) diff --git a/docs/docs/pypaimon/writing.md b/docs/docs/pypaimon/writing.md index ebe524620b7a..f272226b84cf 100644 --- a/docs/docs/pypaimon/writing.md +++ b/docs/docs/pypaimon/writing.md @@ -125,12 +125,17 @@ Reads restore omitted partition columns from split metadata. Selected-schema inputs must include the routing columns they need; `write_arrow(data, bucket=...)` also accepts a bucket assigned upstream. -Changing `with_write_type()` or `with_blob_consumer()` affects writers created -for new partition/bucket pairs. Existing bucket writers retain their original -columns and consumer, including across stream `prepare_commit(identifier)` -calls. Full-schema inputs are projected to each target's captured columns by -name. Selected-schema inputs must match the target's captured layout; missing -columns are rejected before any group from the batch is written. +Configure `with_write_type()` before the first effective write. Empty batches +and ignored events do not lock the selection. After writing, the same ordered +selection can be set again, but changing its columns or their order is rejected, +including after `prepare_commit()` or idle writer cleanup. Create a new writer +to use a different write type. + +Changing `with_blob_consumer()` affects newly created partition/bucket writers. +Active bucket writers retain their consumer across stream +`prepare_commit(identifier)` calls. Unaware bucket writers retire on an idle +prepare; fixed bucket writers also require a newer committed identifier from +the same commit user. A reopened bucket uses the current consumer. Batch `prepare_commit()` accepts no arguments and can be called once. Stream `prepare_commit(identifier)` requires an integer-like value in the signed diff --git a/paimon-python/pypaimon/execution/README.md b/paimon-python/pypaimon/execution/README.md index 699c9b3dad01..17507d6e6ce7 100644 --- a/paimon-python/pypaimon/execution/README.md +++ b/paimon-python/pypaimon/execution/README.md @@ -207,14 +207,17 @@ Invalid call arguments fail before either executor prepares files. Public stream write, update and commit entry points normalize integer-like commit identifiers to Java `long` / Rust `i64` before execution starts. Write-column selections reject empty selections and duplicate names before -replacing the current configuration. Like Java, append writers capture their -write columns and Blob consumer per partition/bucket. Reconfiguration affects -future bucket writers; existing buckets retain their configuration across -`prepare_commit` calls. Full-schema inputs project by column name, and batches -are checked against every target's captured layout before staging any group. -Native named-row conversion queries Rust for the target columns using only -routing metadata, so Python does not duplicate bucket state or encode unused -payloads. Index build/drop requests bind partition +replacing the current configuration. Append writers keep one ordered write type +after their first effective write; empty batches and ignored events do not lock +it. Repeating the same selection is a no-op, while a different selection requires +a new writer, including after prepares or idle cleanup. Full-schema inputs are +projected to the selected columns by name. Native named-row conversion encodes +these columns and the required routing/event fields in one batch, leaving unused +payloads untouched. +Blob consumers are captured per partition/bucket; changes apply to newly created +writers, and active writers retain their consumer across prepares. Idle writers +retire at the Java commit boundary, so reopening uses the current consumer. +Index build/drop requests bind partition predicates to the partition layout and reject predicates containing non-partition fields. diff --git a/paimon-python/pypaimon/execution/native/write/native_write.py b/paimon-python/pypaimon/execution/native/write/native_write.py index 7d726c3c7cd1..c88f9c79806e 100644 --- a/paimon-python/pypaimon/execution/native/write/native_write.py +++ b/paimon-python/pypaimon/execution/native/write/native_write.py @@ -107,14 +107,14 @@ def with_write_type(self, write_cols): # schema. It does not configure a PK column-subset writer. return self raise ValueError('Column-subset primary-key writes require the Python executor') - self._native_writer.with_write_type(names) + self._configure_native('with_write_type', names) self._schema = pa.schema([self._table_schema.field(name) for name in names]) self._write_cols = names return self def with_blob_consumer(self, blob_consumer): if blob_consumer is None: - self._configure_native_blob('with_blob_consumer', None) + self._configure_native('with_blob_consumer', None) self._blob_consumer = None return self if not callable(blob_consumer): @@ -125,7 +125,7 @@ def consume(field_name, encoded): descriptor = None if encoded is None else BlobDescriptor.deserialize(encoded) return blob_consumer(field_name, descriptor) - self._configure_native_blob('with_blob_consumer', consume) + self._configure_native('with_blob_consumer', consume) self._blob_consumer = blob_consumer return self @@ -136,14 +136,14 @@ def with_blob_uri_reader_factory(self, uri_reader_factory): # Only core knows whether RowKind filtering accepted any rows or a # postpone writer buffered data. Validate before changing the live # bridge's fallback, including when clearing the application factory. - self._configure_native_blob('with_blob_uri_reader_factory', self._blob_rows) + self._configure_native('with_blob_uri_reader_factory', self._blob_rows) self._blob_rows.fallback = uri_reader_factory else: - self._configure_native_blob('with_blob_uri_reader_factory', uri_reader_factory) + self._configure_native('with_blob_uri_reader_factory', uri_reader_factory) self._blob_uri_reader_factory = uri_reader_factory return self - def _configure_native_blob(self, method, value): + def _configure_native(self, method, value): try: getattr(self._native_writer, method)(value) except ValueError as error: @@ -182,7 +182,7 @@ def _prepare_native_arrow_data(self, data): if len(set(names)) == len(names) and all(name in self._table_schema.names for name in names): expected = pa.schema([self._table_schema.field(name) for name in names]) if arrow_schemas_compatible(data.schema, expected, **compatible): - # Core validates the destination's current or captured layout. + # Core validates the selected write type and routing columns. # Python checks only the declared types of transported fields. return data raise ValueError( @@ -209,41 +209,21 @@ def write_row(self, row): logical_columns = self._write_cols if full_row else list(values) if skip_write_row(self.table, values, row.get_row_kind(), logical_columns): return - if self._write_cols is not None or (self._written and not self.table.is_primary_key_table): + require_columns(values, names, 'write_row') + if self._write_cols is not None: routing = list(self.table.partition_keys) if ((self.table.bucket_mode() == BucketMode.HASH_FIXED and self.table.options.bucket() > 1) or isinstance(self, NativePostponeFixedBucketTableWrite)): routing.extend(self.table.table_schema.bucket_keys) routing = list(dict.fromkeys(routing)) - require_columns(values, routing, 'write_row') - captured = names - if self._written and not self.table.is_primary_key_table: - if routing: - metadata = row_values_to_arrow_table( - values, self.table.table_schema.fields, routing).to_batches()[0] - else: - from pypaimon.utils.arrow_utils import zero_column_batch - metadata = zero_column_batch(1) - captured = self._native_writer._row_write_columns(metadata) - require_columns(values, captured, 'write_row') row_kind_field = self.table.options.options.to_map().get('rowkind.field') if row_kind_field is not None: routing.append(row_kind_field) require_columns(values, routing, 'write_row') - if captured != names: - # A historical named row carries its captured columns plus - # routing metadata. Preserve its declared column order even if - # routing adds no extra fields. Current-type NOT NULL checks - # have already run above; unused payloads need no encoding. - names = list(captured) + [name for name in self._table_schema.names - if name in routing and name not in captured] - elif full_row or not set(routing).issubset(names): - retained = set(captured) | set(routing) - names = self._table_schema.names - # Unwritten payloads remain untouched, including live Blobs. - values = {name: values.get(name) if name in retained else None for name in names} - else: - require_columns(values, names, 'write_row') + # Keep the selected order, followed by routing/event fields. Values + # outside this transport remain untouched, including live Blobs. + names = list(names) + [name for name in self._table_schema.names + if name in routing and name not in names] try: data = (self._blob_rows.to_batch(self.table, [values], names) if self._blob_rows is not None else diff --git a/paimon-python/pypaimon/execution/python/write/file_store_write.py b/paimon-python/pypaimon/execution/python/write/file_store_write.py index 4e687254d7a3..42eb39a37c41 100644 --- a/paimon-python/pypaimon/execution/python/write/file_store_write.py +++ b/paimon-python/pypaimon/execution/python/write/file_store_write.py @@ -58,10 +58,6 @@ def __init__(self, table, commit_user): self.table: FileStoreTable = table self._commit_user = commit_user self.data_writers: Dict[Tuple, DataWriter] = {} - # Like Java's per-bucket writer factory, existing writers retain the - # requested layout even when the factory changes for future buckets. - # DataWriter.write_cols may instead describe optimized physical files. - self._write_columns_by_bucket = {} self._runtime_total_buckets: Dict[Tuple, int] = {} self._last_modified_commit_identifiers = {} self._written_buckets = set() @@ -99,49 +95,22 @@ def write( data: pa.RecordBatch, total_buckets=None, ): - self.validate_write_columns(data.schema.names, [(partition, bucket)]) self._check_runtime_bucket(partition, bucket, total_buckets) key = (partition, bucket) if key not in self.data_writers: self._new_bucket_writer(partition, bucket) writer = self.data_writers[key] - write_columns = self._write_columns_for(partition, bucket) - if write_columns is not None: - data = project_arrow_batch(data, write_columns) + if self.write_cols is not None: + data = project_arrow_batch(data, self.write_cols) writer.write(data) self._written_buckets.add(key) def _new_bucket_writer(self, partition, bucket): - columns = tuple(self.write_cols) if self.write_cols is not None else None writer = self._create_data_writer(partition, bucket, self.options) key = (partition, bucket) self.data_writers[key] = writer - self._write_columns_by_bucket[key] = columns self._ever_written = True - def _write_columns_for(self, partition, bucket): - return self._write_columns_by_bucket.get((partition, bucket), self.write_cols) - - def validate_write_columns(self, input_columns, partition_buckets): - """Check every target layout before any group from a batch is staged.""" - full_columns = list(self.table.field_names) - input_columns = list(input_columns) - if input_columns == full_columns: - return - from pypaimon.table.special_fields import SpecialFields - # Internal PK rows carry their event kind outside the user schema. - if SpecialFields.VALUE_KIND.name not in full_columns: - input_columns = [name for name in input_columns if name != SpecialFields.VALUE_KIND.name] - if input_columns == full_columns: - return - for partition, bucket in partition_buckets: - columns = self._write_columns_for(partition, bucket) - expected = list(columns) if columns is not None else full_columns - if input_columns != expected: - raise ValueError( - 'Input columns %s do not match captured write columns %s for partition %s, bucket %s' - % (input_columns, expected, partition, bucket)) - def write_row( self, partition: Tuple, @@ -151,13 +120,30 @@ def write_row( total_buckets=None, ): key = (partition, bucket) - columns = self._write_columns_for(partition, bucket) + columns = self.write_cols column_names = list(columns) if columns is not None else list(self.table.field_names) # InternalRows carry explicit field names and may arrive reordered. - # Bind values by those names, then write the bucket's captured order. + # Bind values by those names, then write the selected column order. require_columns(values_by_name, column_names, 'write_row') + writer = self.data_writers.get(key) + data = None + if writer is None and self._has_blob_columns(): + # Dedicated writers consume live Blobs themselves. Validate only + # ordinary selected fields before capturing the writer's layout, + # without opening or encoding any selected or unused Blob value. + ordinary_columns = [name for name in column_names + if not is_blob_file_field(self.table.field_dict[name])] + if ordinary_columns: + row_values_to_arrow_table( + values_by_name, self.table.table_schema.fields, ordinary_columns) + elif writer is None or not hasattr(writer, 'write_row'): + data = row_values_to_arrow_table( + values_by_name, self.table.table_schema.fields, column_names) + from pypaimon.write.row_kind import with_row_kind + data = with_row_kind(self.table, data, row) + self._check_runtime_bucket(partition, bucket, total_buckets) - if key not in self.data_writers: + if writer is None: self._new_bucket_writer(partition, bucket) writer = self.data_writers[key] if hasattr(writer, 'write_row'): @@ -165,13 +151,6 @@ def write_row( self._written_buckets.add(key) return - data = row_values_to_arrow_table( - values_by_name, - self.table.table_schema.fields, - column_names, - ) - from pypaimon.write.row_kind import with_row_kind - data = with_row_kind(self.table, data, row) writer.write(data.to_batches()[0]) self._written_buckets.add(key) @@ -466,7 +445,6 @@ def _release_prepared_files(self): self._last_modified_commit_identifiers[key] = identifier for key in retired: del self.data_writers[key] - self._write_columns_by_bucket.pop(key, None) self._last_modified_commit_identifiers.pop(key, None) partitions = {partition for partition, _ in self.data_writers} for partition in set(self._runtime_total_buckets).difference(partitions): @@ -483,7 +461,6 @@ def close(self): for writer in self.data_writers.values(): writer.close() self.data_writers.clear() - self._write_columns_by_bucket.clear() self._runtime_total_buckets.clear() self._last_modified_commit_identifiers.clear() self._written_buckets.clear() @@ -497,7 +474,6 @@ def abort(self): except Exception as e: logger.warning("Failed to abort data writer.", exc_info=e) self.data_writers.clear() - self._write_columns_by_bucket.clear() self._runtime_total_buckets.clear() self._last_modified_commit_identifiers.clear() self._written_buckets.clear() diff --git a/paimon-python/pypaimon/execution/python/write/table_write.py b/paimon-python/pypaimon/execution/python/write/table_write.py index c3d4e1016557..4a6f9afd21b0 100644 --- a/paimon-python/pypaimon/execution/python/write/table_write.py +++ b/paimon-python/pypaimon/execution/python/write/table_write.py @@ -99,8 +99,6 @@ def write_arrow_batch(self, data: pa.RecordBatch, bucket: Optional[int] = None): return groups = writer.row_key_extractor.extract_partition_bucket_groups(data) - writer.file_store_write.validate_write_columns( - data.schema.names, [(partition, bucket) for partition, bucket, _ in groups]) for partition, bucket, row_indices in groups: if row_indices is None: # Every input row belongs to the same partition/bucket. Passing the @@ -220,6 +218,10 @@ def with_write_type(self, write_cols: List[str]): and set(write_cols) == set(full_names)) if write_cols == full_names or full_pk_type: write_cols = None + if writer.file_store_write._ever_written: + if write_cols == writer.file_store_write.write_cols: + return self + raise RuntimeError('with_write_type must be called before any write operation') writer.file_store_write.write_cols = write_cols return self @@ -358,15 +360,6 @@ def _validate_pyarrow_schema(self, data_schema: pa.Schema, *, check_nested_nulla data_schema, write_cols_schema, check_nested_nullability=check_nested_nullability): return - # Existing bucket writers keep their input layout across checkpoints. - # Their requested columns can differ from optimized physical writeCols. - if isinstance(writer.file_store_write, FileStoreWrite): - for columns in writer.file_store_write._write_columns_by_bucket.values(): - if columns is not None and writer._is_compatible_pyarrow_schema( - data_schema, writer._write_cols_pyarrow_schema(columns), - check_nested_nullability=check_nested_nullability): - return - writer._raise_inconsistent_schema(data_schema) def _is_compatible_pyarrow_schema( diff --git a/paimon-python/pypaimon/tests/native_filtered_writer_configuration_test.py b/paimon-python/pypaimon/tests/native_filtered_writer_configuration_test.py index 417676b59da7..9d3f5d974c33 100644 --- a/paimon-python/pypaimon/tests/native_filtered_writer_configuration_test.py +++ b/paimon-python/pypaimon/tests/native_filtered_writer_configuration_test.py @@ -162,10 +162,11 @@ def test_filtered_blob_batches_allow_configuration_and_real_writes_capture_it(re writer.write_arrow(pa.Table.from_pylist([ dict(id=2, op='+I', payload=BlobDescriptor('custom://source', 6, 7).serialize()), ], schema=schema)) - # Java keeps the existing bucket's columns/consumer while allowing - # a new factory configuration for future buckets. + # Consumers can change for future buckets; the writer's ordered + # column selection remains fixed after an accepted write. writer.with_blob_consumer(None) - writer.with_write_type(['id']) + with pytest.raises(RuntimeError, match='before.*writ'): + writer.with_write_type(['id']) with pytest.raises(RuntimeError, match='before.*writ'): writer.with_blob_uri_reader_factory(None) if native: diff --git a/paimon-python/pypaimon/tests/write/append_row_kind_test.py b/paimon-python/pypaimon/tests/write/append_row_kind_test.py index b6204e13d409..8c9abd7776a0 100644 --- a/paimon-python/pypaimon/tests/write/append_row_kind_test.py +++ b/paimon-python/pypaimon/tests/write/append_row_kind_test.py @@ -363,7 +363,6 @@ def column(self, index): writer.table = Mock(field_names=schema.names) writer.write_cols = columns writer.data_writers = {} - writer._write_columns_by_bucket = {} writer._runtime_total_buckets = {} writer._written_buckets = set() writer._ever_written = False diff --git a/paimon-python/pypaimon/tests/write/append_writer_lifecycle_test.py b/paimon-python/pypaimon/tests/write/append_writer_lifecycle_test.py index ea0bbceb3816..7c6f1cffc127 100644 --- a/paimon-python/pypaimon/tests/write/append_writer_lifecycle_test.py +++ b/paimon-python/pypaimon/tests/write/append_writer_lifecycle_test.py @@ -34,71 +34,81 @@ pytest.mark.python_write, pytest.mark.python_commit] +def _blob_table(tmp_path, native): + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.create_database('default', True) + fields = [DataField(0, 'id', AtomicType('INT')), DataField(1, 'payload', AtomicType('BLOB')), + DataField(2, 'pt', AtomicType('STRING'))] + options = {'write.native.enabled': str(native).lower(), 'commit.native.enabled': 'false', + 'scan.native-plan.enabled': 'false', 'read.native.enabled': 'false', + 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'} + catalog.create_table('default.blobs', Schema(fields=fields, partition_keys=['pt'], options=options), False) + return catalog.get_table('default.blobs') + + +def _blob_rows(table): + read = table.new_read_builder() + return sorted(read.new_read().to_arrow(read.new_scan().plan().splits()).to_pylist(), + key=lambda row: row['id']) + + @pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) @pytest.mark.parametrize('idle_prepares', [0, 1, 2]) @pytest.mark.parametrize('mode', ['arrow', 'row']) -def test_unaware_idle_prepare_reopens_with_current_write_type(tmp_path, native, idle_prepares, mode): - table = _table(tmp_path, native) +def test_unaware_idle_prepare_reopens_with_current_blob_consumer(tmp_path, native, idle_prepares, mode): + table = _blob_table(tmp_path, native) builder = table.new_stream_write_builder() writer, commit = builder.new_write(), builder.new_commit() + received = [] + rows = [{'id': 1, 'payload': b'old', 'pt': 'old'}, {'id': 2, 'payload': b'new', 'pt': 'old'}] try: assert writer.prepare_commit(0) == [] - writer.with_write_type(['id', 'pt']) - _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}], mode) + writer.with_blob_consumer(lambda name, descriptor: received.append(('old', descriptor)) or True) + _write(writer, table, rows[:1], mode) commit.commit(writer.prepare_commit(1), 1) for identifier in range(2, 2 + idle_prepares): assert writer.prepare_commit(identifier) == [] - writer.with_write_type(['value', 'pt']) - _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}], mode) + writer.with_blob_consumer(lambda name, descriptor: received.append(('new', descriptor)) or True) + _write(writer, table, rows[1:], mode) identifier = 2 + idle_prepares - messages = writer.prepare_commit(identifier) - assert [file.write_cols for message in messages for file in message.new_files] == [ - ['value', 'pt'] if idle_prepares else ['id', 'pt']] - commit.commit(messages, identifier) + commit.commit(writer.prepare_commit(identifier), identifier) + assert [label for label, _ in received] == ['old', 'new' if idle_prepares else 'old'] + assert [descriptor.length for _, descriptor in received] == [3, 3] finally: writer.close() commit.close() - expected = [ - {'id': None, 'value': 20, 'pt': 'old'}, {'id': 1, 'value': None, 'pt': 'old'} - ] if idle_prepares else [ - {'id': 1, 'value': None, 'pt': 'old'}, {'id': 2, 'value': None, 'pt': 'old'}] - assert _read(table) == expected + assert _blob_rows(table) == rows @pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) @pytest.mark.parametrize('first_identifier', [-(1 << 63), 0], ids=['minimum', 'ordinary']) def test_minimum_identifier_remains_first_prepare_until_a_later_nonempty_increment( tmp_path, native, first_identifier): - table = _table(tmp_path, native) + table = _blob_table(tmp_path, native) builder = table.new_stream_write_builder() writer, commit = builder.new_write(), builder.new_commit() + received = [] + rows = [dict(id=index, payload=value, pt='old') + for index, value in enumerate([b'one', b'two', b'three'], 1)] try: - writer.with_write_type(['id', 'pt']) - _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}]) + writer.with_blob_consumer(lambda name, descriptor: received.append('old') or True) + _write(writer, table, rows[:1]) commit.commit(writer.prepare_commit(first_identifier), first_identifier) assert writer.prepare_commit(first_identifier + 1) == [] - writer.with_write_type(['value', 'pt']) - _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}]) - messages = writer.prepare_commit(first_identifier + 1) - minimum = first_identifier == -(1 << 63) - assert [file.write_cols for message in messages for file in message.new_files] == [ - ['id', 'pt'] if minimum else ['value', 'pt']] - commit.commit(messages, first_identifier + 1) - # Once a nonempty increment has advanced past MIN, this idle bucket - # can retire and capture the factory's new full layout. + writer.with_blob_consumer(lambda name, descriptor: received.append('second') or True) + _write(writer, table, rows[1:2]) + commit.commit(writer.prepare_commit(first_identifier + 1), first_identifier + 1) + assert received == ['old', 'old' if first_identifier == -(1 << 63) else 'second'] + # A later nonempty increment advances past MIN and permits retirement. assert writer.prepare_commit(first_identifier + 2) == [] - writer.with_write_type(table.field_names) - _write(writer, table, [{'id': 3, 'value': 30, 'pt': 'old'}]) - messages = writer.prepare_commit(first_identifier + 3) - assert [file.write_cols for message in messages for file in message.new_files] == [None] - commit.commit(messages, first_identifier + 3) + writer.with_blob_consumer(lambda name, descriptor: received.append('third') or True) + _write(writer, table, rows[2:]) + commit.commit(writer.prepare_commit(first_identifier + 3), first_identifier + 3) + assert received[-1] == 'third' finally: writer.close() commit.close() - expected = [{'id': 1, 'value': None, 'pt': 'old'}, {'id': 3, 'value': 30, 'pt': 'old'}] - expected.append({'id': 2, 'value': None, 'pt': 'old'} if minimum - else {'id': None, 'value': 20, 'pt': 'old'}) - assert _read(table) == sorted(expected, key=lambda row: row['id'] or -1) + assert _blob_rows(table) == rows @pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) @@ -174,8 +184,12 @@ def test_fixed_idle_prepare_requires_strictly_newer_commit_of_same_user(tmp_path writer.with_write_type(['id', 'pt']) _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}], 'precomputed') commit.commit(writer.prepare_commit(1), 1) + if not native: + reference = weakref.ref(writer.file_store_write.data_writers[(('old',), 0)]) # Equality is insufficient, even if another checkpoint is being prepared. assert writer.prepare_commit(2) == [] + if not native: + assert reference() is not None expected = [{'id': 1, 'value': None, 'pt': 'old'}] if latest_commit == 'other-user': other = table.new_stream_write_builder() @@ -192,18 +206,16 @@ def test_fixed_idle_prepare_requires_strictly_newer_commit_of_same_user(tmp_path commit.commit(writer.prepare_commit(2), 2) expected.insert(0, {'id': 8, 'value': None, 'pt': 'helper'}) assert writer.prepare_commit(3) == [] - writer.with_write_type(['value', 'pt']) + if not native: + gc.collect() + assert (reference() is None) == (latest_commit == 'higher-id') + assert writer.with_write_type(['id', 'pt']) is writer _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}], 'precomputed') - messages = writer.prepare_commit(4) - retired = latest_commit == 'higher-id' - assert [file.write_cols for message in messages for file in message.new_files] == [ - ['value', 'pt'] if retired else ['id', 'pt']] - commit.commit(messages, 4) + commit.commit(writer.prepare_commit(4), 4) finally: writer.close() commit.close() - expected.append({'id': None, 'value': 20, 'pt': 'old'} if retired - else {'id': 2, 'value': None, 'pt': 'old'}) + expected.append({'id': 2, 'value': None, 'pt': 'old'}) assert _read(table) == sorted(expected, key=lambda row: (row['pt'], row['id'] or -1)) @@ -271,7 +283,7 @@ def test_fixed_snapshot_search_stops_at_expired_gap(tmp_path, missing, lookup): side_effect=AssertionError('writer cleanup must not use the catalog snapshot loader')): assert writer.prepare_commit(2) == [] assert get.call_args_list == ([] if lookup == 'latest' else [((2,), {})]) - writer.with_write_type(['value', 'pt']) + assert (('old',), 0) in writer.file_store_write.data_writers _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}], 'precomputed') messages = writer.prepare_commit(3) assert [file.write_cols for message in messages for file in message.new_files] == [['id', 'pt']] @@ -371,7 +383,7 @@ def test_retirement_close_failure_keeps_active_files_for_retry(tmp_path): messages = writer.prepare_commit(2) assert [message.partition for message in messages] == [('new',)] assert [file.row_count for message in messages for file in message.new_files] == [1] - assert set(store.data_writers) == set(store._write_columns_by_bucket) == {(('new',), 0)} + assert set(store.data_writers) == {(('new',), 0)} files = list(tmp_path.rglob('*.parquet')) assert len(files) == 2 writer.abort() diff --git a/paimon-python/pypaimon/tests/write/write_factory_reconfiguration_test.py b/paimon-python/pypaimon/tests/write/write_factory_reconfiguration_test.py index 3a354601ea9e..142256c958e5 100644 --- a/paimon-python/pypaimon/tests/write/write_factory_reconfiguration_test.py +++ b/paimon-python/pypaimon/tests/write/write_factory_reconfiguration_test.py @@ -15,8 +15,9 @@ # specific language governing permissions and limitations # under the License. -"""Writer-factory changes apply to future buckets, as in Java.""" +"""A writer uses one ordered write type after accepting its first row.""" +import io from unittest.mock import Mock, patch import pyarrow as pa @@ -32,17 +33,21 @@ not native_runtime_available(), reason='pypaimon-rust not installed')])] -def _table(tmp_path, native, precomputed=False, unwritten_not_null=False): +def _table(tmp_path, native, precomputed=False, unwritten_not_null=False, blob=False): catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) catalog.create_database('default', True) fields = [pa.field('id', pa.int32()), pa.field('value', pa.int32())] if unwritten_not_null: fields.append(pa.field('omitted', pa.int32(), nullable=False)) + if blob: + fields.append(pa.field('payload', pa.large_binary())) schema = pa.schema(fields + [pa.field('pt', pa.string())]) options = {'write.native.enabled': str(native).lower(), 'commit.native.enabled': 'false', 'scan.native-plan.enabled': 'false', 'read.native.enabled': 'false'} if precomputed: options.update({'bucket': '2', 'bucket-key': 'id'}) + if blob: + options.update({'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}) catalog.create_table('default.factory', Schema.from_pyarrow_schema( schema, partition_keys=['pt'], options=options), False) return catalog.get_table('default.factory') @@ -71,96 +76,119 @@ def _read(table): @pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) @pytest.mark.parametrize('mode', ['arrow', 'row', 'precomputed']) +@pytest.mark.parametrize('stage', ['write', 'checkpoint', 'idle']) @pytest.mark.parametrize('initial_full', [False, True], ids=['subset', 'full']) -def test_existing_bucket_keeps_its_write_type_across_factory_changes_and_prepare( - tmp_path, native, mode, initial_full): +def test_write_type_is_fixed_after_first_write_even_after_idle_retirement( + tmp_path, native, mode, stage, initial_full): table = _table(tmp_path, native, mode == 'precomputed') builder = table.new_stream_write_builder() writer, commit = builder.new_write(), builder.new_commit() + columns = table.field_names if initial_full else ['id', 'pt'] try: if not initial_full: - writer.with_write_type(['id', 'pt']) + writer.with_write_type(columns) _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}], mode) - commit.commit(writer.prepare_commit(1), 1) - writer.with_write_type(['value', 'pt']) + if stage != 'write': + commit.commit(writer.prepare_commit(1), 1) + if stage == 'idle': + assert writer.prepare_commit(2) == [] + assert writer.with_write_type(list(columns)) is writer + # Column order is part of the type, even when the field set is equal. + for different in (list(reversed(columns)), ['value', 'pt']): + with pytest.raises(RuntimeError, match='before any write'): + writer.with_write_type(different) _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}, {'id': 3, 'value': 30, 'pt': 'new'}], mode) - messages = writer.prepare_commit(2) - assert {message.partition: [file.write_cols for file in message.new_files] - for message in messages} == { - ('old',): [None if initial_full else ['id', 'pt']], - ('new',): [['value', 'pt']], - } - commit.commit(messages, 2) + messages = writer.prepare_commit(3) + assert all(file.write_cols == (None if initial_full else columns) + for message in messages for file in message.new_files) + commit.commit(messages, 3) finally: writer.close() commit.close() assert _read(table) == [ - {'id': None, 'value': 30, 'pt': 'new'}, + {'id': 3, 'value': 30 if initial_full else None, 'pt': 'new'}, {'id': 1, 'value': 10 if initial_full else None, 'pt': 'old'}, {'id': 2, 'value': 20 if initial_full else None, 'pt': 'old'}, ] @pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) -@pytest.mark.parametrize('mode', ['arrow', 'row']) -def test_historical_subset_is_valid_only_for_its_existing_bucket(tmp_path, native, mode): - table = _table(tmp_path, native) +@pytest.mark.parametrize('input_', ['row', 'batch', 'table', 'pandas']) +@pytest.mark.parametrize('kind', [None, '-D', '-U'], ids=['empty', 'delete', 'update-before']) +def test_empty_or_ignored_input_does_not_fix_write_type(tmp_path, native, input_, kind): + from pypaimon.tests.write.append_row_kind_test import _table as event_table, _write as write_event + table, schema = event_table(tmp_path, native, **{'ignore-delete': 'true'}) builder = table.new_batch_write_builder() writer, commit = builder.new_write(), builder.new_commit() try: writer.with_write_type(['id', 'pt']) - _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}], mode) - writer.with_write_type(['value', 'pt']) - _write(writer, table, [{'id': 2, 'pt': 'old'}], mode, ['id', 'pt']) - _write(writer, table, [{'value': 30, 'pt': 'new'}], mode, ['value', 'pt']) + values = [] if kind is None else [dict(id=1, op=kind, pt='ignored')] + if kind is None and input_ == 'batch': + writer.write_arrow_batch(pa.RecordBatch.from_arrays( + [pa.array([], type=field.type) for field in schema], schema=schema)) + else: + write_event(writer, table, schema, values, input_, True) + assert writer.with_write_type(['op', 'pt']) is writer + write_event(writer, table, schema, [dict(id=2, op='+I', pt='kept')], input_, True) + assert writer.with_write_type(['op', 'pt']) is writer + with pytest.raises(RuntimeError, match='before any write'): + writer.with_write_type(['id', 'pt']) commit.commit(writer.prepare_commit()) finally: writer.close() commit.close() + read = table.new_read_builder() + assert read.new_read().to_arrow(read.new_scan().plan().splits()).to_pylist() == [ + dict(id=None, op='+I', pt='kept')] + + +@pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) +@pytest.mark.parametrize('mode', ['arrow', 'row']) +def test_different_write_types_use_different_writers(tmp_path, native, mode): + table = _table(tmp_path, native) + builder = table.new_batch_write_builder() + left = builder.new_write().with_write_type(['id', 'pt']) + right = builder.new_write().with_write_type(['value', 'pt']) + commit = builder.new_commit() + try: + _write(left, table, [{'id': 1, 'value': 10, 'pt': 'old'}], mode) + _write(right, table, [{'id': 2, 'value': 20, 'pt': 'old'}], mode) + commit.commit(left.prepare_commit() + right.prepare_commit()) + finally: + left.close() + right.close() + commit.close() assert _read(table) == [ - {'id': None, 'value': 30, 'pt': 'new'}, - {'id': 1, 'value': None, 'pt': 'old'}, - {'id': 2, 'value': None, 'pt': 'old'}, - ] + {'id': None, 'value': 20, 'pt': 'old'}, {'id': 1, 'value': None, 'pt': 'old'}] @pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) -@pytest.mark.parametrize('transport', ['current', 'historical']) -def test_incompatible_mixed_batch_is_rejected_before_any_group_is_staged(tmp_path, native, transport): +def test_incompatible_batch_is_rejected_before_any_partition_is_staged(tmp_path, native): table = _table(tmp_path, native) builder = table.new_batch_write_builder() writer, commit = builder.new_write(), builder.new_commit() try: writer.with_write_type(['id', 'pt']) _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}]) - writer.with_write_type(['value', 'pt']) - if transport == 'current': - rows, columns = [{'value': 20, 'pt': 'new'}, {'value': 30, 'pt': 'old'}], ['value', 'pt'] - else: - rows, columns = [{'id': 2, 'pt': 'old'}, {'id': 3, 'pt': 'new'}], ['id', 'pt'] before = set(tmp_path.rglob('*.parquet')) - with pytest.raises((ValueError, RuntimeError)): - _write(writer, table, rows, columns=columns) + with pytest.raises(ValueError): + _write(writer, table, [{'value': 20, 'pt': 'new'}, {'value': 30, 'pt': 'old'}], + columns=['value', 'pt']) assert set(tmp_path.rglob('*.parquet')) == before - if not native: - assert set(writer.file_store_write.data_writers) == {(('old',), 0)} - # A failed batch must retain earlier pending rows and admit a valid one. - _write(writer, table, [{'id': 4, 'value': 40, 'pt': 'old'}, - {'id': 5, 'value': 50, 'pt': 'new'}]) + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}, + {'id': 3, 'value': 30, 'pt': 'new'}]) commit.commit(writer.prepare_commit()) finally: writer.close() commit.close() assert _read(table) == [ - {'id': None, 'value': 50, 'pt': 'new'}, - {'id': 1, 'value': None, 'pt': 'old'}, - {'id': 4, 'value': None, 'pt': 'old'}, - ] + {'id': 3, 'value': None, 'pt': 'new'}, + {'id': 1, 'value': None, 'pt': 'old'}, {'id': 2, 'value': None, 'pt': 'old'}] @pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) -def test_reordered_full_named_rows_keep_the_old_bucket_column_names(tmp_path, native): +def test_reordered_named_rows_keep_the_selected_column_order(tmp_path, native): table = _table(tmp_path, native, unwritten_not_null=True) builder = table.new_batch_write_builder() writer, commit = builder.new_write(), builder.new_commit() @@ -168,44 +196,18 @@ def test_reordered_full_named_rows_keep_the_old_bucket_column_names(tmp_path, na writer.with_write_type(['pt', 'id']) _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old', 'omitted': None}], 'row', ['pt', 'value', 'id', 'omitted']) - writer.with_write_type(['pt', 'value']) - _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old', 'omitted': None}, - {'id': 3, 'value': 30, 'pt': 'new', 'omitted': None}], + _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old', 'omitted': None}], 'row', ['value', 'id', 'pt', 'omitted']) commit.commit(writer.prepare_commit()) finally: writer.close() commit.close() assert _read(table) == [ - {'id': None, 'value': 30, 'pt': 'new', 'omitted': None}, {'id': 1, 'value': None, 'pt': 'old', 'omitted': None}, - {'id': 2, 'value': None, 'pt': 'old', 'omitted': None}, - ] - - -@pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) -def test_row_missing_an_old_bucket_column_cannot_reinterpret_current_factory_values(tmp_path, native): - table = _table(tmp_path, native) - builder = table.new_batch_write_builder() - writer, commit = builder.new_write(), builder.new_commit() - try: - writer.with_write_type(['id', 'pt']) - _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}], 'row') - writer.with_write_type(['value', 'pt']) - with pytest.raises((ValueError, RuntimeError)): - _write(writer, table, [{'value': 99, 'pt': 'old'}], 'row', ['value', 'pt']) - _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}], 'row') - commit.commit(writer.prepare_commit()) - finally: - writer.close() - commit.close() - assert _read(table) == [ - {'id': 1, 'value': None, 'pt': 'old'}, - {'id': 2, 'value': None, 'pt': 'old'}, - ] + {'id': 2, 'value': None, 'pt': 'old', 'omitted': None}] -def test_failed_bucket_creation_does_not_capture_a_factory_layout(tmp_path): +def test_failed_bucket_creation_does_not_fix_write_type(tmp_path): table = _table(tmp_path, False) writer = table.new_batch_write_builder().new_write().with_write_type(['value', 'pt']) try: @@ -213,7 +215,7 @@ def test_failed_bucket_creation_does_not_capture_a_factory_layout(tmp_path): with patch.object(store, '_create_data_writer', side_effect=OSError('cannot create writer')): with pytest.raises(OSError, match='cannot create writer'): _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}]) - assert store.data_writers == {} and store._write_columns_by_bucket == {} + assert store.data_writers == {} writer.with_write_type(['id', 'pt']) _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}]) messages = writer.prepare_commit() @@ -230,11 +232,117 @@ def test_requested_write_columns_are_distinct_from_optimized_physical_metadata(t optimized = Mock(write_cols=['id']) with patch.object(store, '_create_data_writer', return_value=optimized): _write(writer, table, [{'id': 1, 'value': 10, 'pt': 'old'}]) - writer.with_write_type(['pt', 'value']) + assert writer.with_write_type(['pt', 'id']) is writer _write(writer, table, [{'id': 2, 'value': 20, 'pt': 'old'}]) - assert store._write_columns_by_bucket[(('old',), 0)] == ('pt', 'id') assert [call[0][0].to_pydict() for call in optimized.write.call_args_list] == [ {'pt': ['old'], 'id': [1]}, {'pt': ['old'], 'id': [2]}] store.data_writers.clear() finally: writer.close() + + +@pytest.mark.parametrize('backend', ['python', 'native-mock']) +@pytest.mark.parametrize('blob_input', ['none', 'omitted', 'selected']) +def test_first_row_encoding_failure_does_not_fix_write_type(tmp_path, backend, blob_input): + from pypaimon.execution.native.write.native_write import NativeTableWrite as NativeWrite + from pypaimon.table.row.blob import Blob + + table = _table(tmp_path, False, blob=blob_input != 'none') + builder = table.new_batch_write_builder() + core = Mock(spec_set=['with_blob_uri_reader_factory', 'with_write_type', 'write_arrow', 'close']) + writer = builder.new_write() if backend == 'python' else NativeWrite(table, 'job', None, False, core) + payload = Mock(spec=Blob) + payload.new_input_stream.side_effect = AssertionError('unused Blob was opened') + payload.to_descriptor.side_effect = AssertionError('unused Blob was encoded') + columns = ['id', 'pt'] if blob_input != 'selected' else ['id', 'payload', 'pt'] + values = {'id': 'invalid-int', 'pt': 'old'} + if blob_input != 'none': + values['payload'] = payload + try: + writer.with_write_type(columns) + with pytest.raises(pa.ArrowInvalid): + writer.write_row(GenericRow(list(values.values()), [table.field_dict[name] for name in values])) + assert not list(tmp_path.rglob('*.parquet')) and not list(tmp_path.rglob('*.blob')) + if backend == 'python': + assert writer.file_store_write.data_writers == {} + assert not writer.file_store_write._ever_written + else: + core.write_arrow.assert_not_called() + assert not writer._written + assert writer.with_write_type(['value', 'pt']) is writer + values['id'] = 1 + values['value'] = 20 + writer.write_row(GenericRow(list(values.values()), [table.field_dict[name] for name in values])) + if backend == 'python': + commit = builder.new_commit() + try: + commit.commit(writer.prepare_commit()) + finally: + commit.close() + expected = {'id': None, 'value': 20, 'pt': 'old'} + if blob_input != 'none': + expected['payload'] = None + assert _read(table) == [expected] + else: + core.write_arrow.assert_called_once() + assert core.write_arrow.call_args[0][0].to_pydict() == {'value': [20], 'pt': ['old']} + payload.new_input_stream.assert_not_called() + payload.to_descriptor.assert_not_called() + finally: + writer.close() + + +def test_blob_row_encoding_retry_consumes_the_selected_live_blob_once(tmp_path): + from pypaimon.table.row.blob import Blob + + class InputBlob(Blob): + opened = 0 + + def to_data(self): + raise AssertionError('live Blob was materialized') + + def to_descriptor(self): + raise AssertionError('live Blob was encoded') + + def new_input_stream(self): + self.opened += 1 + return io.BytesIO(b'payload') + + table = _table(tmp_path, False, blob=True) + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + payload = InputBlob() + try: + writer.with_write_type(['id', 'payload', 'pt']) + with pytest.raises(pa.ArrowInvalid): + writer.write_row(GenericRow(['invalid-int', payload, 'old'], + [table.field_dict[name] for name in ['id', 'payload', 'pt']])) + assert payload.opened == 0 + writer.with_write_type(['value', 'payload', 'pt']) + writer.write_row(GenericRow([20, payload, 'old'], + [table.field_dict[name] for name in ['value', 'payload', 'pt']])) + assert payload.opened == 1 + commit.commit(writer.prepare_commit()) + assert payload.opened == 1 + finally: + writer.close() + commit.close() + assert _read(table) == [dict(id=None, value=20, payload=b'payload', pt='old')] + + +def test_failure_after_row_encoding_keeps_the_created_write_type_fixed(tmp_path): + from pypaimon.execution.python.write.writer.append_only_data_writer import AppendOnlyDataWriter + + table = _table(tmp_path, False) + writer = table.new_batch_write_builder().new_write().with_write_type(['id', 'pt']) + failure = OSError('writer failed after row encoding') + try: + with patch.object(AppendOnlyDataWriter, 'write', side_effect=failure): + with pytest.raises(OSError) as caught: + writer.write_row(GenericRow([1, 'old'], [table.field_dict[name] for name in ['id', 'pt']])) + assert caught.value is failure + assert writer.with_write_type(['id', 'pt']) is writer + with pytest.raises(RuntimeError, match='before any write'): + writer.with_write_type(['value', 'pt']) + finally: + writer.close() diff --git a/paimon-python/pypaimon/tests/write/writer_execution_protocol_test.py b/paimon-python/pypaimon/tests/write/writer_execution_protocol_test.py index 4b988e603112..3cb081571a9e 100644 --- a/paimon-python/pypaimon/tests/write/writer_execution_protocol_test.py +++ b/paimon-python/pypaimon/tests/write/writer_execution_protocol_test.py @@ -184,28 +184,26 @@ def _named_row(table, values, kind=None): return row -def _captured_writer(table): - core = Mock() +def _selected_writer(table, columns=None): + core = Mock(spec_set=['with_blob_uri_reader_factory', 'with_write_type', 'write_arrow', 'close']) writer = NativeWrite(table, 'rows', None, True, core) - writer.with_write_type(['id']) - writer.write_row(_named_row(table, {'id': 111, 'pt': 1})) - writer.with_write_type(['value']) - core._row_write_columns.return_value = ['id'] + columns = columns if columns is not None else ['id'] + writer.with_write_type(columns) + values = {'id': 111, 'pt': 1, 'value': 222} + writer.write_row(_named_row(table, {name: values[name] for name in dict.fromkeys(columns + ['pt'])})) core.write_arrow.reset_mock() return writer, core -def test_native_full_named_row_retains_captured_columns_and_leaves_unused_blob_untouched(row_table): +def test_native_full_named_row_encodes_selected_columns_and_leaves_unused_blob_untouched(row_table): from pypaimon.table.row.blob import Blob blob = Mock(spec=Blob) blob.new_input_stream.side_effect = AssertionError('unwritten Blob was opened') blob.to_descriptor.side_effect = AssertionError('unwritten Blob was encoded') - writer, core = _captured_writer(row_table) + writer, core = _selected_writer(row_table) try: writer.write_row(_named_row(row_table, {'unused': blob, 'value': 333, 'pt': 1, 'id': 222})) - metadata = core._row_write_columns.call_args[0][0] - assert metadata.schema.names == ['pt'] - assert metadata.to_pydict() == {'pt': [1]} + core.write_arrow.assert_called_once() payload = core.write_arrow.call_args[0][0] assert payload.schema.names == ['id', 'pt'] assert payload.to_pydict() == {'id': [222], 'pt': [1]} @@ -216,11 +214,9 @@ def test_native_full_named_row_retains_captured_columns_and_leaves_unused_blob_u writer.close() -def test_native_historical_partial_named_row_omits_current_non_null_payload(row_table): - writer, core = _captured_writer(row_table) +def test_native_partial_named_row_omits_unselected_non_null_payload(row_table): + writer, core = _selected_writer(row_table) try: - # This row uses the old bucket's id layout. The new factory's NOT NULL - # value is absent, and must not be manufactured as an invalid NULL. writer.write_row(_named_row(row_table, {'id': 222, 'pt': 1})) payload = core.write_arrow.call_args[0][0] assert payload.schema.names == ['id', 'pt'] @@ -229,55 +225,46 @@ def test_native_historical_partial_named_row_omits_current_non_null_payload(row_ writer.close() -@pytest.mark.parametrize('selection', [['value', 'unused'], ['pt', 'id', 'value', 'unused']]) -def test_native_reconfigured_full_row_encodes_only_destination_payloads(row_table, selection): - from pypaimon.table.row.blob import Blob - blob = Mock(spec=Blob) - blob.new_input_stream.side_effect = AssertionError('unwritten Blob was opened') - blob.to_descriptor.side_effect = AssertionError('unwritten Blob was encoded') - writer, core = _captured_writer(row_table) +def test_native_type_change_rejected_by_core_preserves_adapter_configuration(row_table): + writer, core = _selected_writer(row_table) try: - writer.with_write_type(selection) - writer.write_row(_named_row(row_table, {'pt': 1, 'id': 222, 'value': 333, 'unused': blob})) - payload = core.write_arrow.call_args[0][0] - assert payload.schema.names == ['id', 'pt'] - assert payload.to_pydict() == {'id': [222], 'pt': [1]} - blob.new_input_stream.assert_not_called() - blob.to_descriptor.assert_not_called() + failure = ValueError('with_write_type must be called before any write operation') + core.with_write_type.side_effect = failure + with pytest.raises(RuntimeError, match='before any write') as caught: + writer.with_write_type(['value']) + assert caught.value.__cause__ is failure + assert writer._schema.names == writer._write_cols == ['id'] + writer.write_row(_named_row(row_table, {'pt': 1, 'id': 222})) + assert core.write_arrow.call_args[0][0].to_pydict() == {'id': [222], 'pt': [1]} finally: writer.close() -@pytest.mark.parametrize('full_current_type', [False, True]) -def test_native_historical_row_preserves_captured_column_order(row_table, full_current_type): - core = Mock() - writer = NativeWrite(row_table, 'ordered', None, True, core) +@pytest.mark.parametrize('columns', [['id', 'pt'], ['pt', 'id']]) +def test_native_named_row_preserves_selected_column_order(row_table, columns): + writer, core = _selected_writer(row_table, columns) try: - writer.with_write_type(['id', 'pt']) - writer.write_row(_named_row(row_table, {'pt': 1, 'id': 111})) - writer.with_write_type(row_table.field_names if full_current_type else ['value']) - core._row_write_columns.return_value = ['id', 'pt'] writer.write_row(_named_row(row_table, {'pt': 1, 'id': 222})) payload = core.write_arrow.call_args[0][0] - assert payload.schema.names == ['id', 'pt'] + assert payload.schema.names == columns assert payload.to_pydict() == {'id': [222], 'pt': [1]} finally: writer.close() -def test_native_full_row_current_not_null_validation_precedes_captured_column_lookup(row_table): - writer, core = _captured_writer(row_table) +def test_native_selected_not_null_validation_precedes_encoding(row_table): + writer, core = _selected_writer(row_table, ['value']) try: - with pytest.raises(ValueError, match='value'): - writer.write_row(_named_row(row_table, {'pt': 1, 'id': 222, 'value': None, 'unused': None})) - core._row_write_columns.assert_not_called() + with patch.object(writer._blob_rows, 'to_batch', side_effect=AssertionError('payload was encoded')): + with pytest.raises(ValueError, match='value'): + writer.write_row(_named_row(row_table, {'pt': 1, 'id': 222, 'value': None, 'unused': None})) core.write_arrow.assert_not_called() finally: writer.close() -def test_native_row_missing_captured_value_fails_before_encoding_or_staging(row_table): - writer, core = _captured_writer(row_table) +def test_native_row_missing_selected_value_fails_before_encoding_or_staging(row_table): + writer, core = _selected_writer(row_table) try: with patch.object(writer._blob_rows, 'to_batch', side_effect=AssertionError('payload was encoded')): with pytest.raises(ValueError, match='id'): @@ -287,31 +274,26 @@ def test_native_row_missing_captured_value_fails_before_encoding_or_staging(row_ writer.close() -def test_native_ignored_row_does_not_query_routing_or_encode_payloads(row_table): +def test_native_ignored_row_does_not_encode_or_stage_payloads(row_table): from pypaimon.table.row.row_kind import RowKind - writer, core = _captured_writer(row_table) + writer, core = _selected_writer(row_table) try: - core._row_write_columns.reset_mock() with patch.object(writer._blob_rows, 'to_batch', side_effect=AssertionError('ignored payload was encoded')): writer.write_row(_named_row(row_table, {'pt': 1, 'id': 222, 'value': 333}, RowKind.UPDATE_BEFORE)) - core._row_write_columns.assert_not_called() core.write_arrow.assert_not_called() finally: writer.close() -def test_native_unpartitioned_row_queries_a_zero_column_metadata_row(tmp_path): +def test_native_unpartitioned_rows_only_call_the_write_interface(tmp_path): table = _table(tmp_path) - core = Mock() - core._row_write_columns.return_value = ['id'] + core = Mock(spec_set=['with_write_type', 'write_arrow', 'close']) writer = NativeWrite(table, 'rows', None, True, core) try: writer.with_write_type(['id']) writer.write_row(_named_row(table, {'id': 1})) writer.write_row(_named_row(table, {'id': 2})) - metadata = core._row_write_columns.call_args[0][0] - assert metadata.num_rows == 1 - assert metadata.num_columns == 0 - assert core.write_arrow.call_args[0][0].column('id').to_pylist() == [2] + assert core.write_arrow.call_count == 2 + assert core.write_arrow.call_args[0][0].to_pydict() == {'id': [2]} finally: writer.close() From f5c4d31d65e812c78a12d7160383d06297c2f67a Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 22 Oct 2026 12:11:01 +0800 Subject: [PATCH 9/9] Align append input encoding and routing transport boundaries --- docs/docs/pypaimon/writing.md | 8 +- paimon-python/pypaimon/execution/README.md | 9 +- .../python/write/file_store_write.py | 22 +- .../execution/python/write/table_write.py | 53 ++++- .../write_factory_reconfiguration_test.py | 99 ++++++++- .../tests/write/write_transport_input_test.py | 209 ++++++++++++++++++ 6 files changed, 382 insertions(+), 18 deletions(-) create mode 100644 paimon-python/pypaimon/tests/write/write_transport_input_test.py diff --git a/docs/docs/pypaimon/writing.md b/docs/docs/pypaimon/writing.md index f272226b84cf..8deec81f9bcb 100644 --- a/docs/docs/pypaimon/writing.md +++ b/docs/docs/pypaimon/writing.md @@ -119,14 +119,16 @@ remain unsupported. Append `with_write_type()` preserves the selected column order. Both the full table schema and the selected schema are accepted as Arrow input; only selected -columns are written. Full Arrow inputs and named Rows can carry partition and -hash columns used only for routing, independently of the physical write type. +columns are written. Selected Arrow inputs can also carry partition, hash and +configured row-kind columns used for routing and event handling. The selected +columns must retain their declared order. Named Rows bind fields by name. Reads restore omitted partition columns from split metadata. Selected-schema inputs must include the routing columns they need; `write_arrow(data, bucket=...)` also accepts a bucket assigned upstream. Configure `with_write_type()` before the first effective write. Empty batches -and ignored events do not lock the selection. After writing, the same ordered +and ignored events do not lock the selection. Row-to-Arrow encoding failures +before execution starts also leave it configurable. After writing, the same ordered selection can be set again, but changing its columns or their order is rejected, including after `prepare_commit()` or idle writer cleanup. Create a new writer to use a different write type. diff --git a/paimon-python/pypaimon/execution/README.md b/paimon-python/pypaimon/execution/README.md index 17507d6e6ce7..ddc9cc06cdf6 100644 --- a/paimon-python/pypaimon/execution/README.md +++ b/paimon-python/pypaimon/execution/README.md @@ -211,9 +211,12 @@ replacing the current configuration. Append writers keep one ordered write type after their first effective write; empty batches and ignored events do not lock it. Repeating the same selection is a no-op, while a different selection requires a new writer, including after prepares or idle cleanup. Full-schema inputs are -projected to the selected columns by name. Native named-row conversion encodes -these columns and the required routing/event fields in one batch, leaving unused -payloads untouched. +projected to the selected columns by name. Selected Arrow inputs may also carry +routing/event fields while retaining the selected column order. Named-row +conversion encodes selected ordinary and inline reference fields before capturing +a bucket writer; encoding failures leave the configuration open. Dedicated live +Blob payloads are consumed only during execution, and unused payloads remain +untouched. Blob consumers are captured per partition/bucket; changes apply to newly created writers, and active writers retain their consumer across prepares. Idle writers retire at the Java commit boundary, so reopening uses the current consumer. diff --git a/paimon-python/pypaimon/execution/python/write/file_store_write.py b/paimon-python/pypaimon/execution/python/write/file_store_write.py index 42eb39a37c41..21758d212c6b 100644 --- a/paimon-python/pypaimon/execution/python/write/file_store_write.py +++ b/paimon-python/pypaimon/execution/python/write/file_store_write.py @@ -24,7 +24,7 @@ from pypaimon.common.options.core_options import CoreOptions from pypaimon.schema.data_types import is_blob_file_field from pypaimon.write.commit_message import CommitMessage -from pypaimon.write.row_utils import require_columns, row_values_to_arrow_table +from pypaimon.write.row_utils import inline_blob_value, require_columns, row_values_to_arrow_table from pypaimon.schema.arrow_schema import project_arrow_batch from pypaimon.execution.python.write.writer.append_only_data_writer import AppendOnlyDataWriter from pypaimon.execution.python.write.writer.dedicated_format_writer import DedicatedFormatWriter @@ -128,11 +128,23 @@ def write_row( writer = self.data_writers.get(key) data = None if writer is None and self._has_blob_columns(): - # Dedicated writers consume live Blobs themselves. Validate only - # ordinary selected fields before capturing the writer's layout, - # without opening or encoding any selected or unused Blob value. + # Dedicated writers consume live payloads themselves. Inline fields + # instead encode to Arrow before the writer captures its layout. + descriptor_fields = self.options.blob_descriptor_fields() + view_fields = self.options.blob_view_fields() + inline_columns = [name for name in column_names if name in descriptor_fields or name in view_fields] + if inline_columns: + values_by_name = dict(values_by_name) + for name in inline_columns: + values_by_name[name] = inline_blob_value( + values_by_name[name], name in descriptor_fields, name in view_fields) + # Reuse the encoded bytes in the dedicated row path. Keep the + # original event kind and leave unselected Blob objects alone. + from pypaimon.table.row.generic_row import GenericRow + row = GenericRow([values_by_name[name] for name in column_names], + [self.table.field_dict[name] for name in column_names], row.get_row_kind()) ordinary_columns = [name for name in column_names - if not is_blob_file_field(self.table.field_dict[name])] + if name in inline_columns or not is_blob_file_field(self.table.field_dict[name])] if ordinary_columns: row_values_to_arrow_table( values_by_name, self.table.table_schema.fields, ordinary_columns) diff --git a/paimon-python/pypaimon/execution/python/write/table_write.py b/paimon-python/pypaimon/execution/python/write/table_write.py index 4a6f9afd21b0..917a050e0ddd 100644 --- a/paimon-python/pypaimon/execution/python/write/table_write.py +++ b/paimon-python/pypaimon/execution/python/write/table_write.py @@ -21,7 +21,7 @@ from pypaimon.common.options.core_options import ChangelogProducer from pypaimon.schema.arrow_schema import ( - arrow_schemas_compatible, normalize_arrow_strings, normalize_arrow_write_nullability) + arrow_schemas_compatible, normalize_arrow_strings, normalize_arrow_write_nullability, project_arrow_batch) from pypaimon.schema.data_types import PyarrowFieldParser from pypaimon.table.row.blob import BlobConsumer from pypaimon.write.row_utils import ( @@ -303,21 +303,51 @@ def _prepare_arrow_data(self, data): return data data = writer._normalize_arrow_write_schema(data) from pypaimon.write.row_kind import filter_write_batch - write_cols = (writer.file_store_write.write_cols - if data.schema.names == writer.table_pyarrow_schema.names else data.schema.names) + selected = (writer.file_store_write.write_cols if not table.is_primary_key_table else None) + write_cols = selected + if selected is None: + write_cols = (writer.file_store_write.write_cols + if data.schema.names == writer.table_pyarrow_schema.names else data.schema.names) data = filter_write_batch(table, data, write_cols) if not data.num_rows: return data + if selected is not None: + routing = self._write_routing_columns() + columns = selected + [name for name in data.schema.names + if name in routing and name not in selected] + data = self._project_write_input(data, columns) + # Routing fields are not written values. Arrow's whole-table + # string cast must not reapply their declared top-level NOT NULL. + schema = pa.schema([field if field.name in selected else field.with_nullable(True) + for field in data.schema], metadata=data.schema.metadata) + data = (pa.RecordBatch.from_arrays(data.columns, schema=schema) if isinstance(data, pa.RecordBatch) + else pa.Table.from_arrays(data.columns, schema=schema)) data = normalize_arrow_strings(data) + physical = self._project_write_input(data, selected) if selected is not None else data target = writer.table_pyarrow_schema - if data.schema.names != target.names: - target = writer._write_cols_pyarrow_schema(data.schema.names) + if physical.schema.names != target.names: + target = writer._write_cols_pyarrow_schema(physical.schema.names) # Validate only surviving values. Keep the input transport metadata: # PyArrow otherwise rejects hidden NULL children of NULL ROW parents. - validated = normalize_arrow_write_nullability(data, target) + validated = normalize_arrow_write_nullability(physical, target) writer._validate_pyarrow_schema(validated.schema) return data + @staticmethod + def _project_write_input(data, columns): + return (project_arrow_batch(data, columns) if isinstance(data, pa.RecordBatch) + else data.select(columns)) + + def _write_routing_columns(self): + table = self.writer.table + routing = set(table.partition_keys) + if table.options.bucket() > 1: + routing.update(table.table_schema.bucket_keys) + row_kind = table.options.options.to_map().get('rowkind.field') + if row_kind is not None: + routing.add(row_kind) + return routing + def _prepare_arrow_schema(self, data): # Legacy schema-only callers have no row-kind options to evaluate. # A real writer must defer casts until ignored events have been removed. @@ -359,6 +389,17 @@ def _validate_pyarrow_schema(self, data_schema: pa.Schema, *, check_nested_nulla if writer._is_compatible_pyarrow_schema( data_schema, write_cols_schema, check_nested_nullability=check_nested_nullability): return + table = getattr(writer, 'table', None) + if table is not None and not table.is_primary_key_table: + names = data_schema.names + selected = set(write_cols) + allowed = selected | self._write_routing_columns() + if (len(set(names)) == len(names) and set(names).issubset(allowed) + and [name for name in names if name in selected] == write_cols): + transport_schema = writer._write_cols_pyarrow_schema(names) + if writer._is_compatible_pyarrow_schema( + data_schema, transport_schema, check_nested_nullability=check_nested_nullability): + return writer._raise_inconsistent_schema(data_schema) diff --git a/paimon-python/pypaimon/tests/write/write_factory_reconfiguration_test.py b/paimon-python/pypaimon/tests/write/write_factory_reconfiguration_test.py index 142256c958e5..161f689f8706 100644 --- a/paimon-python/pypaimon/tests/write/write_factory_reconfiguration_test.py +++ b/paimon-python/pypaimon/tests/write/write_factory_reconfiguration_test.py @@ -33,7 +33,7 @@ not native_runtime_available(), reason='pypaimon-rust not installed')])] -def _table(tmp_path, native, precomputed=False, unwritten_not_null=False, blob=False): +def _table(tmp_path, native, precomputed=False, unwritten_not_null=False, blob=False, unwritten_blob=False): catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) catalog.create_database('default', True) fields = [pa.field('id', pa.int32()), pa.field('value', pa.int32())] @@ -41,6 +41,8 @@ def _table(tmp_path, native, precomputed=False, unwritten_not_null=False, blob=F fields.append(pa.field('omitted', pa.int32(), nullable=False)) if blob: fields.append(pa.field('payload', pa.large_binary())) + if unwritten_blob: + fields.append(pa.field('unused_blob', pa.large_binary())) schema = pa.schema(fields + [pa.field('pt', pa.string())]) options = {'write.native.enabled': str(native).lower(), 'commit.native.enabled': 'false', 'scan.native-plan.enabled': 'false', 'read.native.enabled': 'false'} @@ -346,3 +348,98 @@ def test_failure_after_row_encoding_keeps_the_created_write_type_fixed(tmp_path) writer.with_write_type(['value', 'pt']) finally: writer.close() + + +@pytest.mark.parametrize('backend', ['python', 'native-mock']) +@pytest.mark.parametrize('inline', ['descriptor', 'view']) +def test_inline_blob_conversion_failure_does_not_fix_write_type(tmp_path, backend, inline): + from pypaimon.execution.native.write.native_write import NativeTableWrite as NativeWrite + from pypaimon.table.row.blob import Blob, BlobView + + table = _table(tmp_path, False, blob=True).copy({'blob-%s-field' % inline: 'payload'}) + builder = table.new_batch_write_builder() + core = Mock(spec_set=['with_blob_uri_reader_factory', 'with_write_type', 'write_arrow', 'close']) + writer = builder.new_write() if backend == 'python' else NativeWrite(table, 'job', None, False, core) + value = Blob.from_data(b'not a descriptor') + if inline == 'view': + value = Mock(spec=BlobView) + value.view_struct.serialize.side_effect = ValueError('cannot serialize view') + value.new_input_stream.side_effect = AssertionError('view stream was opened') + try: + writer.with_write_type(['payload', 'pt']) + with pytest.raises(ValueError): + writer.write_row(GenericRow([value, 'old'], [table.field_dict[name] for name in ['payload', 'pt']])) + assert not list(tmp_path.rglob('*.parquet')) and not list(tmp_path.rglob('*.blob')) + if backend == 'python': + assert writer.file_store_write.data_writers == {} + assert not writer.file_store_write._ever_written + else: + core.write_arrow.assert_not_called() + assert not writer._written + assert writer.with_write_type(['value', 'pt']) is writer + writer.write_row(GenericRow([20, 'old'], [table.field_dict[name] for name in ['value', 'pt']])) + if backend == 'python': + commit = builder.new_commit() + try: + commit.commit(writer.prepare_commit()) + finally: + commit.close() + assert _read(table) == [dict(id=None, value=20, payload=None, pt='old')] + else: + assert core.write_arrow.call_args[0][0].to_pydict() == {'value': [20], 'pt': ['old']} + finally: + writer.close() + + +@pytest.mark.parametrize('inline', ['descriptor', 'view']) +def test_inline_blob_wire_validation_failure_keeps_write_type_fixed(tmp_path, inline): + table = _table(tmp_path, False, blob=True).copy({'blob-%s-field' % inline: 'payload'}) + writer = table.new_batch_write_builder().new_write().with_write_type(['payload', 'pt']) + try: + with pytest.raises(ValueError, match='blob-%s-field' % inline): + writer.write_row(GenericRow([b'invalid wire bytes', 'old'], + [table.field_dict[name] for name in ['payload', 'pt']])) + assert writer.file_store_write._ever_written + assert writer.with_write_type(['payload', 'pt']) is writer + with pytest.raises(RuntimeError, match='before any write'): + writer.with_write_type(['value', 'pt']) + finally: + writer.close() + + +@pytest.mark.parametrize('inline', ['descriptor', 'view']) +def test_inline_blob_encoding_is_reused_and_unselected_blob_is_untouched(tmp_path, inline): + import pyarrow.parquet as pq + from pypaimon.table.row.blob import Blob, BlobDescriptor, BlobView, BlobViewStruct + + table = _table(tmp_path, False, blob=True, unwritten_blob=True).copy({'blob-%s-field' % inline: 'payload'}) + if inline == 'descriptor': + descriptor = BlobDescriptor('file:///unused-source', 0, 7) + encoded = descriptor.serialize() + value = Mock(spec=Blob) + value.to_descriptor.side_effect = [descriptor, AssertionError('descriptor encoded twice')] + else: + encoded = BlobViewStruct('default.upstream', 1, 2).serialize() + value = Mock(spec=BlobView) + value.view_struct.serialize.side_effect = [encoded, AssertionError('view encoded twice')] + value.new_input_stream.side_effect = AssertionError('inline Blob stream was opened') + unused = Mock(spec=Blob) + unused.to_descriptor.side_effect = AssertionError('unselected Blob was encoded') + unused.new_input_stream.side_effect = AssertionError('unselected Blob stream was opened') + writer = table.new_batch_write_builder().new_write().with_write_type(['payload', 'pt']) + try: + names = ['payload', 'unused_blob', 'pt'] + writer.write_row(GenericRow([value, unused, 'old'], [table.field_dict[name] for name in names])) + messages = writer.prepare_commit() + paths = [file.file_path for message in messages for file in message.new_files] + assert len(paths) == 1 + assert pq.ParquetFile(paths[0]).read(columns=['payload']).column('payload').to_pylist() == [encoded] + if inline == 'descriptor': + value.to_descriptor.assert_called_once() + else: + value.view_struct.serialize.assert_called_once() + value.new_input_stream.assert_not_called() + unused.to_descriptor.assert_not_called() + unused.new_input_stream.assert_not_called() + finally: + writer.close() diff --git a/paimon-python/pypaimon/tests/write/write_transport_input_test.py b/paimon-python/pypaimon/tests/write/write_transport_input_test.py new file mode 100644 index 000000000000..80adb0c5d633 --- /dev/null +++ b/paimon-python/pypaimon/tests/write/write_transport_input_test.py @@ -0,0 +1,209 @@ +# 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. + +"""Selected Arrow values can carry separate routing and event columns.""" + +import pyarrow as pa +import pytest + +from pypaimon import CatalogFactory, Schema +from pypaimon.execution.native.runtime import native_runtime_available + +pytestmark = [pytest.mark.python_plan, pytest.mark.python_read, + pytest.mark.python_write, pytest.mark.python_commit] +_NATIVE = [False, pytest.param(True, marks=[pytest.mark.native_plan, pytest.mark.skipif( + not native_runtime_available(), reason='pypaimon-rust not installed')])] + + +def _table(tmp_path, native, bucket=-1, generated=False): + fields = pa.schema([ + pa.field('id', pa.int32(), nullable=False), + pa.field('value', pa.int32(), nullable=False), + ('other', pa.int32()), + ('unused', pa.struct([pa.field('item', pa.string(), nullable=False)])), + pa.field('pt', pa.string(), nullable=False), + pa.field('op', pa.string(), nullable=False), + ]) + options = {'write.native.enabled': str(native).lower(), 'commit.native.enabled': 'false', + 'read.native.enabled': 'false', 'scan.native-plan.enabled': 'false', + 'bucket': str(bucket)} + if bucket > 0: + options['bucket-key'] = 'id' + else: + options.update({'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}) + if generated: + options.update({'rowkind.field': 'op', 'ignore-delete': 'true'}) + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.create_database('default', True) + catalog.create_table('default.transport', Schema.from_pyarrow_schema( + fields, partition_keys=['pt'], options=options), False) + return catalog.get_table('default.transport'), fields + + +def _batch(schema, columns, rows): + return pa.Table.from_pylist(rows, schema=pa.schema([schema.field(name) for name in columns])).to_batches()[0] + + +def _values(table): + read = table.new_read_builder().with_projection(['value', 'other']) + return sorted(read.new_read().to_arrow(read.new_scan().plan().splits()).to_pylist(), + key=lambda row: row['value']) + + +@pytest.mark.parametrize('native', _NATIVE, ids=['python', 'native']) +@pytest.mark.parametrize('input_', ['batch', 'table']) +def test_selected_values_accept_partition_transport(tmp_path, native, input_): + table, schema = _table(tmp_path, native) + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.with_write_type(['value', 'other']) + input_schema = pa.schema([field.with_type(pa.large_string()) if field.name == 'pt' else field + for field in schema]) + batch = _batch(input_schema, ['pt', 'value', 'other'], [{'pt': 'old', 'value': 10, 'other': 20}]) + writer.write_arrow(batch if input_ == 'batch' else pa.Table.from_batches([batch])) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + assert _values(table) == [{'value': 10, 'other': 20}] + + +@pytest.mark.parametrize('native', _NATIVE, ids=['python', 'native']) +@pytest.mark.parametrize('full', [False, True], ids=['transport', 'full-table']) +@pytest.mark.parametrize('generated', [False, True], ids=['insert', 'rowkind']) +def test_unselected_nulls_and_nested_payloads_do_not_validate_as_written_values( + tmp_path, native, full, generated): + table, schema = _table(tmp_path, native, generated=generated) + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + row = {'id': None, 'value': 10, 'other': 20, 'unused': {'item': None}, + 'pt': None, 'op': '+I' if generated else None} + input_schema = pa.schema([ + field.with_type(pa.struct([('item', pa.large_string())])) if field.name == 'unused' else + field.with_type(pa.large_string()) if field.name == 'pt' else field for field in schema]) + columns = schema.names if full else ['value', 'pt', 'other'] + (['op'] if generated else []) + try: + writer.with_write_type(['value', 'other']) + writer.write_arrow(_batch(input_schema, columns, [row])) + messages = writer.prepare_commit() + assert [message.partition for message in messages] == [(None,)] + assert all(file.write_cols == ['value', 'other'] for message in messages for file in message.new_files) + commit.commit(messages) + finally: + writer.close() + commit.close() + assert _values(table) == [{'value': 10, 'other': 20}] + + +@pytest.mark.parametrize('native', _NATIVE, ids=['python', 'native']) +def test_selected_null_is_checked_before_event_filtering(tmp_path, native): + table, schema = _table(tmp_path, native, generated=True) + writer = table.new_batch_write_builder().new_write().with_write_type(['value', 'other']) + try: + with pytest.raises(ValueError, match='Cannot write null to non-null column.*value'): + writer.write_arrow(_batch(schema, ['value', 'other', 'pt', 'op'], [ + {'value': None, 'other': 20, 'pt': 'ignored', 'op': '-D'}])) + assert writer.with_write_type(['other']) is writer + assert not list(tmp_path.rglob('*.parquet')) + finally: + writer.close() + + +@pytest.mark.parametrize('native', _NATIVE, ids=['python', 'native']) +def test_ignored_transport_needs_no_routing_values_or_captured_layout(tmp_path, native): + table, schema = _table(tmp_path, native, generated=True) + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.with_write_type(['value', 'other']) + writer.write_arrow(_batch(schema, ['value', 'other', 'op'], [ + {'value': 10, 'other': 20, 'op': '-D'}])) + assert writer.with_write_type(['value']) is writer + writer.write_arrow(_batch(schema, ['value', 'pt', 'op'], [ + {'value': 30, 'pt': 'keep', 'op': '+I'}])) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + assert _values(table) == [{'value': 30, 'other': None}] + + +@pytest.mark.parametrize('native', _NATIVE, ids=['python', 'native']) +@pytest.mark.parametrize('columns', [ + ['other', 'value', 'pt'], + ['value', 'other', 'pt', 'id'], + ['value', 'other', 'pt', 'op'], + ['value', 'other', 'pt', 'pt'], + ['value', 'other', 'pt', 'unknown'], + ['other', 'pt'], +], ids=['reordered', 'unneeded-hash-key', 'unconfigured-event', 'duplicate', 'unknown', 'missing-selected']) +def test_transport_rejects_payload_extras_and_changed_selected_layout(tmp_path, native, columns): + table, schema = _table(tmp_path, native) + writer = table.new_batch_write_builder().new_write().with_write_type(['value', 'other']) + fields = [pa.field('unknown', pa.int32()) if name == 'unknown' else schema.field(name) + for name in columns] + arrays = [pa.array(['old'] if field.name in ('pt', 'op') else [10], type=field.type) + for field in fields] + try: + with pytest.raises(ValueError): + writer.write_arrow(pa.RecordBatch.from_arrays(arrays, schema=pa.schema(fields))) + assert writer.with_write_type(['other']) is writer + assert not list(tmp_path.rglob('*.parquet')) + finally: + writer.close() + + +@pytest.mark.parametrize('native', _NATIVE, ids=['python', 'native']) +@pytest.mark.parametrize('explicit', [False, True], ids=['automatic', 'explicit']) +def test_fixed_transport_requires_hash_keys_only_for_automatic_assignment(tmp_path, native, explicit): + table, schema = _table(tmp_path, native, bucket=4) + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + row = {'id': 7, 'value': 10, 'other': 20, 'pt': 'old'} + try: + writer.with_write_type(['value', 'other']) + transport = _batch(schema, ['value', 'other', 'pt'], [row]) + if explicit: + writer.write_arrow(transport, bucket=3) + else: + with pytest.raises(ValueError, match='routing column.*id'): + writer.write_arrow(transport) + assert writer.with_write_type(['other']) is writer + writer.with_write_type(['value', 'other']) + writer.write_arrow(_batch(schema, ['pt', 'id', 'value', 'other'], [row])) + messages = writer.prepare_commit() + assert len(messages) == 1 + if explicit: + assert messages[0].bucket == 3 + commit.commit(messages) + finally: + writer.close() + commit.close() + assert _values(table) == [{'value': 10, 'other': 20}] + + +@pytest.mark.parametrize('native', _NATIVE, ids=['python', 'native']) +def test_transport_still_requires_partition_for_surviving_rows(tmp_path, native): + table, schema = _table(tmp_path, native) + writer = table.new_batch_write_builder().new_write().with_write_type(['value', 'other']) + try: + with pytest.raises(ValueError, match='routing column.*pt'): + writer.write_arrow(_batch(schema, ['value', 'other'], [{'value': 10, 'other': 20}])) + assert writer.with_write_type(['other']) is writer + finally: + writer.close()