You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@paimon.apache.org by lz...@apache.org on 2026/10/22 05:16:30 UTC
(paimon) branch master updated: [Python] Delegate Ray search through backend task execution (#10724)
This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 36e2c902e3 [Python] Delegate Ray search through backend task execution (#10724)
36e2c902e3 is described below
commit 36e2c902e338c8c28882cb3a63e2c57a08b3efbb
Author: Jingsong Lee <ji...@gmail.com>
AuthorDate: Thu Oct 22 13:16:23 2026 +0800
[Python] Delegate Ray search through backend task execution (#10724)
---
.github/workflows/ci-python.yml | 3 +-
docs/docs/pypaimon/multimodal-search.md | 63 ++-
docs/docs/pypaimon/writing.md | 23 +-
paimon-python/conftest.py | 2 +
paimon-python/dev/test_packaging.py | 6 +
paimon-python/pypaimon/execution/README.md | 46 +-
.../native/read/native_full_text_search.py | 2 +-
.../execution/native/read/native_hybrid_search.py | 16 +-
.../execution/native/read/native_search_plan.py | 304 +++++++-----
.../execution/native/read/native_vector_search.py | 4 +-
.../execution/native/read/search_execution.py | 107 ++++
.../execution/native/write/native_write.py | 50 +-
.../read/search/distributed_full_text_search.py} | 86 +---
.../python/read/search/distributed_search.py | 63 +++
.../read/search/distributed_vector_search.py | 281 +++++++++++
.../execution/python/read/search/full_text_read.py | 8 +-
.../execution/python/read/search/hybrid_search.py | 46 +-
.../read/search/primary_key_full_text_read.py | 4 +-
.../read/search/primary_key_full_text_scan.py | 17 +-
.../python/read/search/primary_key_vector_read.py | 4 +-
.../python/read/search/primary_key_vector_scan.py | 17 +-
.../python/read/search/vector_search_read.py | 4 +
.../execution/python/write/file_store_write.py | 165 +++++--
.../pypaimon/execution/python/write/table_write.py | 70 ++-
paimon-python/pypaimon/multimodal/query.py | 8 +-
paimon-python/pypaimon/ray/batch_vector_search.py | 137 +-----
paimon-python/pypaimon/ray/full_text_search.py | 150 +-----
paimon-python/pypaimon/ray/search_execution.py | 201 ++++++++
paimon-python/pypaimon/ray/vector_index_build.py | 4 +-
paimon-python/pypaimon/ray/vector_search.py | 269 +---------
.../pypaimon/table/source/hybrid_search_builder.py | 13 +
.../pypaimon/table/source/search_dispatch.py | 43 +-
.../pypaimon/table/source/search_execution.py | 73 +++
.../pypaimon/table/source/search_types.py | 46 +-
paimon-python/pypaimon/tests/blob_table_test.py | 38 +-
.../tests/data_file_path_directory_test.py | 17 +-
.../pypaimon/tests/datasource_ray_adapter_test.py | 2 +-
.../pypaimon/tests/lazy_search_execution_test.py | 543 +++++++++++++++++++--
.../native_filtered_writer_configuration_test.py | 7 +-
.../pypaimon/tests/native_full_text_search_test.py | 11 +-
.../pypaimon/tests/native_hybrid_search_test.py | 4 +-
.../tests/native_index_result_scan_test.py | 10 +-
.../tests/native_pk_full_text_search_test.py | 28 +-
.../pypaimon/tests/native_search_plan_test.py | 157 +++---
.../tests/native_test_instrumentation_test.py | 30 ++
.../pypaimon/tests/native_vector_context_test.py | 11 +-
.../tests/native_vector_index_build_test.py | 4 +-
.../pypaimon/tests/native_vector_search_test.py | 26 +-
paimon-python/pypaimon/tests/native_write_test.py | 4 +-
.../tests/primary_key_index_definitions_test.py | 31 ++
.../tests/ray_batch_vector_refinement_test.py | 27 +-
.../pypaimon/tests/ray_batch_vector_search_test.py | 48 +-
.../pypaimon/tests/ray_full_text_search_test.py | 22 +-
.../pypaimon/tests/ray_hybrid_parallel_test.py | 71 +--
.../pypaimon/tests/ray_search_execution_test.py | 345 +++++++++++++
.../tests/ray_vector_filter_exactness_test.py | 34 +-
.../pypaimon/tests/ray_vector_refinement_test.py | 29 +-
.../pypaimon/tests/ray_vector_search_test.py | 68 +--
.../tests/search_execution_test_helpers.py | 33 ++
.../pypaimon/tests/search_task_execution_test.py | 328 +++++++++++++
.../tests/table_update_arrow_compatibility_test.py | 7 +-
paimon-python/pypaimon/tests/variant_test.py | 2 +-
.../pypaimon/tests/write/append_row_kind_test.py | 3 +-
.../tests/write/append_writer_lifecycle_test.py | 392 +++++++++++++++
.../pypaimon/tests/write/backend_boundary_test.py | 46 ++
.../write/write_factory_reconfiguration_test.py | 359 +++++++++++---
.../tests/write/write_transport_input_test.py | 209 ++++++++
.../tests/write/writer_execution_protocol_test.py | 94 ++--
paimon-python/pypaimon/write/write_dispatch.py | 12 +-
69 files changed, 4027 insertions(+), 1360 deletions(-)
diff --git a/.github/workflows/ci-python.yml b/.github/workflows/ci-python.yml
index b14c3d5424..775c578943 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 4470bf3d4e..57513fae0f 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 ebe524620b..8deec81f9b 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 b36b91d6f0..394fcda7e3 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 pytest_configure(config):
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 8803119e91..89fabc9003 100644
--- a/paimon-python/dev/test_packaging.py
+++ b/paimon-python/dev/test_packaging.py
@@ -122,6 +122,12 @@ class PackagingTest(unittest.TestCase):
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 48c0d05ea6..ddc9cc06cd 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 9ba1505c62..9ce5b5397b 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 b40a7676d2..660c59593a 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.table.file_store_table import FileStoreTable
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 e55f812243..e409920878 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.search_types import (
)
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 @@ class NativeSearchScan:
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 83ea544fc5..3499eb5009 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 0000000000..13241198d1
--- /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 18a08b0d03..c88f9c7980 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 @@ class NativeTableWrite:
# 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 @@ class NativeTableWrite:
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 @@ class NativeTableWrite:
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 @@ class NativeTableWrite:
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/ray/full_text_search.py b/paimon-python/pypaimon/execution/python/read/search/distributed_full_text_search.py
similarity index 58%
copy from paimon-python/pypaimon/ray/full_text_search.py
copy to paimon-python/pypaimon/execution/python/read/search/distributed_full_text_search.py
index d3eb37142e..010a609fd4 100644
--- a/paimon-python/pypaimon/ray/full_text_search.py
+++ b/paimon-python/pypaimon/execution/python/read/search/distributed_full_text_search.py
@@ -15,41 +15,26 @@
# specific language governing permissions and limitations
# under the License.
-"""Ray full-text shard execution, preserving local corpus and merge semantics."""
-
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from contextlib import closing
-from threading import Event
+"""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.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.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
-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()
- 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):
+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._concurrency = concurrency
- self._remote_args = remote_args
+ self._executor = executor
def _worker_reader(self):
reader = DataEvolutionFullTextRead(
@@ -61,19 +46,22 @@ class _RayFullTextRead(DataEvolutionFullTextRead):
def read(self, splits):
# Keep bitmap construction and candidate data reads off the driver.
index_splits, raw_splits = _split_search_splits(splits)
- indexed = self._eval_column_query({self._text_columns[0].name: index_splits}, None)
+ splits_by_column = {}
+ for split in index_splits:
+ splits_by_column.setdefault(split.column_name, []).append(split)
+ indexed = self._eval_column_query(splits_by_column, None)
raw = self._read_raw_search(_raw_row_ranges(raw_splits), _index_type(index_splits))
return indexed.or_(raw).top_k(self._limit)
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)
- 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:
+ 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():
@@ -86,8 +74,8 @@ class _RayFullTextRead(DataEvolutionFullTextRead):
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:
+ 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)
@@ -127,43 +115,3 @@ def _search_raw(reader, item):
candidates = RoaringBitmap64.and_(candidates, include)
include = reader._matching_rows(candidates)
return _scores(reader._read_raw_search(ranges, index_type, include))
-
-
-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)
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 0000000000..783798d358
--- /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 0000000000..78ac0462cc
--- /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 a7ce89a484..57bab789d4 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 @@ class DataEvolutionFullTextRead(FullTextRead):
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 55ead01bd6..b53d7cd231 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 @@ _MAX_ROUTE_WORKERS = 4
_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 5538cd237b..29f8ee55c9 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 @@ class PrimaryKeyFullTextRead(DataEvolutionFullTextRead):
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 290b5e3bee..a08ba3bb10 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 @@ class PrimaryKeyFullTextScan(FullTextScan):
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 2da2ae6ea6..7bdfbdf84a 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 @@ from pypaimon.utils.roaring_bitmap import RoaringBitmap64
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 d2ca03291f..ffe67e5fc0 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 @@ class PrimaryKeyVectorScan(VectorSearchScan):
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 8abc0d6a32..56a4587ad5 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 @@ class AbstractVectorSearchReadImpl:
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 497dd50a95..21758d212c 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 @@ import pyarrow as pa
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 @@ from pypaimon.table.bucket_mode import BucketMode
logger = logging.getLogger(__name__)
+_MIN_COMMIT_IDENTIFIER = -(1 << 63)
class FileStoreWrite:
@@ -55,12 +56,14 @@ class FileStoreWrite:
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 @@ class FileStoreWrite:
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 @@ class FileStoreWrite:
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 @@ class FileStoreWrite:
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 @@ class FileStoreWrite:
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 @@ class FileStoreWrite:
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 @@ class FileStoreWrite:
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 505ef95748..917a050e0d 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 @@ import pyarrow as pa
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 @@ class PythonTableWrite:
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 @@ class PythonTableWrite:
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 @@ class PythonTableWrite:
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 @@ class PythonTableWrite:
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 @@ class PythonTableWrite:
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 8df4287ddc..9d5375fbcf 100644
--- a/paimon-python/pypaimon/multimodal/query.py
+++ b/paimon-python/pypaimon/multimodal/query.py
@@ -607,9 +607,9 @@ class _PreFilterQuery(ScanQuery):
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 @@ class BatchVectorQuery(_PreFilterQuery):
"""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 6e2eed1318..7a5bdc84f4 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 d3eb37142e..82c5a92236 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 0000000000..91a4f1f643
--- /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 b0c4d1c7a2..31bf1a73a5 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 5a326d7946..6400bb2ef1 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 27a1d2fff4..043765dc0c 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 @@ class HybridSearchBuilder(ABC):
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 @@ class HybridSearchBuilderImpl(HybridSearchBuilder):
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 6589a3fe8e..b174450264 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 @@ class _SearchRead:
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 0000000000..6da31b3f95
--- /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 a649be95fc..5e846f67ee 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 @@ class VectorSearchScan(ABC):
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 @@ class VectorSearchRead(ABC):
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 @@ class FullTextScan(ABC):
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 daa02bee19..a972c040de 100755
--- a/paimon-python/pypaimon/tests/blob_table_test.py
+++ b/paimon-python/pypaimon/tests/blob_table_test.py
@@ -6291,13 +6291,17 @@ class BlobConsumerTest(unittest.TestCase):
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 @@ class BlobConsumerTest(unittest.TestCase):
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 80addf9a81..b6e9e4fe5f 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 @@ class DataFilePathDirectoryTest(unittest.TestCase):
"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 @@ class DataFilePathDirectoryTest(unittest.TestCase):
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 @@ class DataFilePathDirectoryTest(unittest.TestCase):
"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 @@ class DataFilePathDirectoryTest(unittest.TestCase):
"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 1cdcc36d2a..7dd6c11b8c 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 test_ray_read_tasks_preserve_overridden_split_grouping_and_metadata():
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 6923756a3b..36313fd9e6 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 417676b59d..9d3f5d974c 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 51d200b626..36cc51c186 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 3eed083379..2ecd100e81 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 33ede91810..ef646e8026 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 ccb2cda388..288600e147 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 19b9bb89b2..2a93bb5565 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 test_pk_native_plans_export_physical_sources_for_python_readers(source, pk_v
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 75bd44b8f9..91c5153a11 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 @@ pytestmark = [pytest.mark.python_plan, pytest.mark.python_read,
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 8163083b55..9c4608a36f 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 38b56791a6..04fd3e5d98 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 af8db8dc39..b7b8c54007 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 test_unevaluable_scalar_filter_scores_indexed_ranges_raw(vector_table, mode)
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 test_scalar_candidates_obey_refinement_switch(rest_catalog, mode, refine, un
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 test_partition_filter_does_not_make_exact_scalar_index_inexact(vector_table,
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 test_scalar_unindexed_ranges_follow_both_search_modes(
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 test_raw_scalar_prefilter_keeps_uncovered_rows_with_overlapping_scalar_file(
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 625cb73f0f..9f0e241d0a 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 b3a5a18470..12ad4339ca 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 @@ from pypaimon.utils.range import Range
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 f449eb3aad..3edb25301f 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 @@ import pytest
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 test_worker_shares_stream_and_keeps_per_query_membership(table, fail_read):
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 93045c56af..c2f04c7162 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 @@ import pytest
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 test_empty_snapshot_does_not_follow_first_commit(table, ray_cluster):
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 test_global_candidates_and_duplicate_precedence(table):
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 c2521a4a3a..42199e809b 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 @@ 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
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 test_shared_scalar_indexes_are_evaluated_once(tmp_path, ray_cluster):
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 755e6a2e53..9c08c8e6a5 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 @@ import pytest
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 test_routes_overlap_within_budget_and_preserve_order(concurrency, count):
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 test_routes_overlap_within_budget_and_preserve_order(concurrency, count):
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 test_route_failure_cancels_sibling_ray_tasks():
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 0000000000..ec337c5506
--- /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 95dddc57d1..ed1d3c02cb 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 test_ray_applies_exact_row_filter_before_worker_top_k(table, ray_cluster, ba
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 test_ray_applies_exact_row_filter_before_worker_top_k(table, ray_cluster, ba
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 386dbaafbc..f984df8fbc 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 @@ import pytest
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 a90b00c820..b6e1c74f5c 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 @@ import pytest
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 test_empty_snapshot_does_not_follow_first_commit(table, ray_cluster):
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 test_warmed_read_context_and_distinct_worker_processes(table, ray_cluster):
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 test_task_retry_keeps_original_snapshot(ordered, table, ray_cluster, tmp_pat
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 test_failure_cancels_outstanding_tasks(ordered, ray_cluster):
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 test_map_tasks_bounds_submissions_and_cancels_on_early_close(ordered, ray_cl
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 test_slow_first_task_bounds_completed_results(ordered, early_close):
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 0000000000..c185bda6f9
--- /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 0000000000..81721a9552
--- /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 8e24201cc0..28e72a054a 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 4690302448..80c3bf9059 100644
--- a/paimon-python/pypaimon/tests/variant_test.py
+++ b/paimon-python/pypaimon/tests/variant_test.py
@@ -1516,7 +1516,7 @@ class TestVariantPaimonTable(unittest.TestCase):
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 e97c494640..8c9abd7776 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 test_partial_batch_projection_uses_columns_and_preserves_schema_metadata(col
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 0000000000..7c6f1cffc1
--- /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 452f84e3c6..7af31c2cf7 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 3a354601ea..161f689f87 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 @@ _ENGINES = [False, pytest.param(True, marks=[pytest.mark.native_plan, pytest.mar
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 0000000000..80adb0c5d6
--- /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 4b988e6031..3cb081571a 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 e7c80dacfb..fb0d2910ca 100644
--- a/paimon-python/pypaimon/write/write_dispatch.py
+++ b/paimon-python/pypaimon/write/write_dispatch.py
@@ -149,23 +149,27 @@ class NativeTableWrite:
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