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/docs/docs/pypaimon/multimodal-search.md b/docs/docs/pypaimon/multimodal-search.md index 4470bf3d4e99..57513fae0f79 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,55 @@ 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)) +``` + +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 + +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. 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. + +Hybrid builders expose `execute(executor=...)` for the same scheduler. +PyPaimon's 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/docs/docs/pypaimon/writing.md b/docs/docs/pypaimon/writing.md index ebe524620b7a..8deec81f9bcb 100644 --- a/docs/docs/pypaimon/writing.md +++ b/docs/docs/pypaimon/writing.md @@ -119,18 +119,25 @@ 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. -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. 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. + +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/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/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..ddc9cc06cdf6 100644 --- a/paimon-python/pypaimon/execution/README.md +++ b/paimon-python/pypaimon/execution/README.md @@ -144,6 +144,30 @@ 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. + +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. +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. + ## Corresponding execution APIs Compare the operation's inputs, results and lifecycle when changing a backend. @@ -183,14 +207,20 @@ 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. 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. +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/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 b40a7676d211..660c59593a5e 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() + else: + from pypaimon.execution.native.read.search_execution import native_task_executor + callback = native_task_executor(table, executor) + try: + result = native._execute_with_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..e40992087844 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,13 @@ # 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,90 +31,108 @@ ) 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): - 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, source_modes): +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 = _NativePkFullTextPlan if full_text else _NativePkVectorPlan - plan = cls(metadata['snapshot_id'] or 0, splits) - plan._snapshot = snapshot - else: - cls = _NativeFullTextPlan if full_text else _NativeVectorPlan - 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 + cls = PrimaryKeyFullTextScanPlan if full_text else PrimaryKeyVectorScanPlan + return cls(0 if snapshot is None else snapshot.id, splits, snapshot) + return (FullTextScanPlan if full_text else VectorSearchScanPlan)(splits, snapshot) + + +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 + + 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 [sdk.IndexFileMeta._deserialize(_index_row(value)) for value in values] + + def ranges(values): + return [(value.from_, value.to) for value in values] + + if isinstance(split, (PrimaryKeyVectorSearchSplit, PrimaryKeyFullTextSearchSplit)): + cls = sdk.PrimaryKeyFullTextSearchSplit if full_text else sdk.PrimaryKeyVectorSearchSplit + 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): + 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 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): + from pypaimon_rust import datafusion as sdk + + 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, captured) for split in plan.splits()], snapshot) + + +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): @@ -129,12 +143,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 +157,116 @@ 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 read_plan(self, plan): - return _result(self._table, self._native.read_plan(plan), self._primary_key) - - def read_batch_plan(self, plan): + 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 + return _plan_to_native(plan, self._full_text) + + def read_plan(self, plan, executor=None): + native = self._get_native_reader() + plan = self._plan(plan) + if executor is None: + 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 = 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 = 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 = native._read_batch_plan_with_executor(plan, callback) + finally: + callback.cancel() + callback.raise_if_failed() return [_result(self._table, result, False) - for result in self._native.read_batch_plan(plan)] - - 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 + 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 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 _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 _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 _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): @@ -193,7 +274,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): @@ -208,7 +289,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): @@ -216,7 +298,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): @@ -224,4 +306,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/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/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/native/write/native_write.py b/paimon-python/pypaimon/execution/native/write/native_write.py index 18a08b0d038d..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,25 +125,25 @@ 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 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. - 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/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..010a609fd4f9 --- /dev/null +++ b/paimon-python/pypaimon/execution/python/read/search/distributed_full_text_search.py @@ -0,0 +1,117 @@ +# 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 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 +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) + 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) + + 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_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(): + 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_task_results(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..783798d3584a --- /dev/null +++ b/paimon-python/pypaimon/execution/python/read/search/distributed_search.py @@ -0,0 +1,63 @@ +# 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.""" + +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 ( + 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..78ac0462cc16 --- /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 +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.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 + + +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_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) + # 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_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) + 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_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) + # 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_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(): + _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_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(): + _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_full_text_scan.py b/paimon-python/pypaimon/execution/python/read/search/primary_key_full_text_scan.py index 290b5e3bee30..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.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_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/primary_key_vector_scan.py b/paimon-python/pypaimon/execution/python/read/search/primary_key_vector_scan.py index d2ca03291f1b..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.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/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/execution/python/write/file_store_write.py b/paimon-python/pypaimon/execution/python/write/file_store_write.py index 497dd50a9531..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 @@ -36,6 +36,7 @@ logger = logging.getLogger(__name__) +_MIN_COMMIT_IDENTIFIER = -(1 << 63) class FileStoreWrite: @@ -55,12 +56,14 @@ 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 @@ -92,46 +95,21 @@ 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 - - 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)) + self._ever_written = True def write_row( self, @@ -142,27 +120,51 @@ 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 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 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) + 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'): writer.write_row(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) def roll_before_group_if_needed(self, row_count: int): for writer in self.data_writers.values(): @@ -374,11 +376,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 +395,74 @@ 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._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. @@ -404,8 +473,10 @@ 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() + self._prepared_writer_lifecycle = None def abort(self): """Abort all data writers and clean up files produced by this write.""" @@ -415,8 +486,10 @@ 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() + 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..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 ( @@ -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 @@ -234,7 +236,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." @@ -301,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. @@ -357,15 +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 - - # 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 + 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/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..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,27 @@ 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): - 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) + def read_plan(self, plan, executor=None): + 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): - 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) + def read_batch_plan(self, plan, executor=None): + 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): - """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 +132,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..5e846f67ee9f 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: @@ -125,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 @@ -139,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/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 6923756a3b6b..36313fd9e671 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(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)) @@ -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,397 @@ 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 + + +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 + 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 + + 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, + 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) + restored = _plan_from_native(search_table, native, full_text, False) + 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 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 + assert restored_files[0].external_path == index.external_path + assert restored_index.scalar_index_files == [scalar] + 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') + + +@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() + + +@pytest.mark.parametrize('full_text', [False, True]) +@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 + 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=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 + 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 + 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',) + 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() + + +@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_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/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_search_plan_test.py b/paimon-python/pypaimon/tests/native_search_plan_test.py index 19b9bb89b2e8..2a93bb5565d6 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 @@ -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': @@ -327,53 +347,39 @@ 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']) @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): @@ -384,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 @@ -424,8 +432,35 @@ 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 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/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_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')), \ 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/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/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..81721a955220 --- /dev/null +++ b/paimon-python/pypaimon/tests/search_task_execution_test.py @@ -0,0 +1,328 @@ +# 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 + + +@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: plan): + yield + + +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_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) + + 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(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_with_executor=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_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 + + +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() + + +@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 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/append_row_kind_test.py b/paimon-python/pypaimon/tests/write/append_row_kind_test.py index e97c4946407f..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,8 +363,9 @@ 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 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..7c6f1cffc127 --- /dev/null +++ b/paimon-python/pypaimon/tests/write/append_writer_lifecycle_test.py @@ -0,0 +1,392 @@ +# 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] + + +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_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_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_blob_consumer(lambda name, descriptor: received.append(('new', descriptor)) or True) + _write(writer, table, rows[1:], mode) + identifier = 2 + idle_prepares + 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() + 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 = _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_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_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_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() + assert _blob_rows(table) == rows + + +@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) + 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() + 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) == [] + 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') + commit.commit(writer.prepare_commit(4), 4) + finally: + writer.close() + commit.close() + expected.append({'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,), {})]) + 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']] + 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) == {(('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() 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/tests/write/write_factory_reconfiguration_test.py b/paimon-python/pypaimon/tests/write/write_factory_reconfiguration_test.py index 3a354601ea9e..161f689f8706 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,23 @@ 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, 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())] if unwritten_not_null: 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'} 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 +78,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 +198,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 +217,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 +234,212 @@ 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() + + +@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() 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() 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