diff --git a/.github/workflows/ci-python.yml b/.github/workflows/ci-python.yml index 1dd9742be001..580a6a2a78a8 100644 --- a/.github/workflows/ci-python.yml +++ b/.github/workflows/ci-python.yml @@ -239,6 +239,7 @@ jobs: - name: Run Python tests with native planning, reading, writing, updating and committing shell: bash env: + DAFT_RUNNER: 'native' PYTEST_ADDOPTS: '--durations=20' PYPAIMON_TEST_NATIVE_PLAN: '1' PYPAIMON_TEST_NATIVE_READ: '1' diff --git a/paimon-python/pypaimon/execution/native/write/native_commit.py b/paimon-python/pypaimon/execution/native/write/native_commit.py index adc819e1aa43..9cf739cf59eb 100644 --- a/paimon-python/pypaimon/execution/native/write/native_commit.py +++ b/paimon-python/pypaimon/execution/native/write/native_commit.py @@ -204,7 +204,10 @@ def from_native_commit_messages(table, messages): def native_commit_error(error): """Keep retry classification structured; never infer it from error text.""" - from pypaimon_rust import RowIdExistenceConflict as NativeRowIdConflict + try: + from pypaimon_rust import RowIdExistenceConflict as NativeRowIdConflict + except ImportError: + return error from pypaimon.write.row_id_conflict import RowIdExistenceConflict if not isinstance(error, NativeRowIdConflict): diff --git a/paimon-python/pypaimon/execution/python/globalindex/vindex/vindex_vector_global_index_reader.py b/paimon-python/pypaimon/execution/python/globalindex/vindex/vindex_vector_global_index_reader.py index ee97f685796d..7a36decc1649 100644 --- a/paimon-python/pypaimon/execution/python/globalindex/vindex/vindex_vector_global_index_reader.py +++ b/paimon-python/pypaimon/execution/python/globalindex/vindex/vindex_vector_global_index_reader.py @@ -59,12 +59,8 @@ def _run_search(self, vectors, limit, include_row_ids, query_options): params = _search_params(self._search_params_type, options, effective_k) filter_bytes = _filter_bytes(include_row_ids) - if n == 1: - ids, distances = self._reader.search( - queries[0], params, filter_bytes=filter_bytes) - return [_result_from_scores(ids, distances, self._metadata.metric)] - if hasattr(self._reader, "search_batch"): + # The batch ABI pads short results; the scalar ABI requires exactly K. batch_result = self._reader.search_batch( np.ascontiguousarray(queries), params, diff --git a/paimon-python/pypaimon/read/read_dispatch.py b/paimon-python/pypaimon/read/read_dispatch.py index a78356f2b637..f91f46c57bda 100644 --- a/paimon-python/pypaimon/read/read_dispatch.py +++ b/paimon-python/pypaimon/read/read_dispatch.py @@ -223,6 +223,9 @@ def plan_files(scanner): native_supported = _native_query_auth_supported(scan.table) else: native_supported = False + native_auth = native_supported and scan.table.options.query_auth_enabled + if not native_auth and scan._TableScan__auth_query(all_columns=True) is not None: + raise TableNoPermissionException(scan.table.identifier) if native_supported: from pypaimon.execution.native.read.plan import NativeTableScan result = NativeTableScan(scan).plan_files(scanner.skip_level0, scanner.only_read_real_buckets) @@ -231,7 +234,7 @@ def plan_files(scanner): scanner._scanned_snapshot = None scanner._scanned_snapshot_id = snapshot_id return entries - if scan._TableScan__auth_query(all_columns=True) is not None: + if native_auth and scan._TableScan__auth_query(all_columns=True) is not None: raise TableNoPermissionException(scan.table.identifier) from pypaimon.execution.python.read.table_scan import PythonTableScan # Let the selected executor refresh cached rules for this unrestricted scope. @@ -253,11 +256,14 @@ def plan_all(self) -> Plan: native_supported = _native_query_auth_supported(self.table) and self._native_plan_supported() else: native_supported = False + native_auth = native_supported and self.table.options.query_auth_enabled + if not native_auth and self._TableScan__auth_query(all_columns=True) is not None: + raise TableNoPermissionException(self.table.identifier) if native_supported: plan = self._try_native_plan(all_files=True) if plan is not None: return plan - if self._TableScan__auth_query(all_columns=True) is not None: + if native_auth and self._TableScan__auth_query(all_columns=True) is not None: raise TableNoPermissionException(self.table.identifier) from pypaimon.execution.python.read.table_scan import PythonTableScan return PythonTableScan(self).plan_all() diff --git a/paimon-python/pypaimon/table/source/search_dispatch.py b/paimon-python/pypaimon/table/source/search_dispatch.py index 963d96ab18c2..c1e2cad4faef 100644 --- a/paimon-python/pypaimon/table/source/search_dispatch.py +++ b/paimon-python/pypaimon/table/source/search_dispatch.py @@ -21,6 +21,28 @@ transported plans. An execution error never triggers another search. """ +from functools import partial + + +def _python_vector_scan(execution): + from pypaimon.execution.python.read.search.factories import new_vector_scan + return new_vector_scan(execution) + + +def _python_vector_read(builder, batch): + from pypaimon.execution.python.read.search.factories import new_batch_vector_read, new_vector_read + return new_batch_vector_read(builder) if batch else new_vector_read(builder) + + +def _python_full_text_scan(execution): + from pypaimon.execution.python.read.search.factories import new_full_text_scan + return new_full_text_scan(execution) + + +def _python_full_text_read(execution): + from pypaimon.execution.python.read.search.factories import new_full_text_read + return new_full_text_read(execution) + class _SearchScan: def __init__(self, native, python_factory): @@ -78,10 +100,7 @@ def _native_requested(builder): def create_vector_scan(builder, execution, batch=False): - def python_factory(): - from pypaimon.execution.python.read.search.factories import new_vector_scan - return new_vector_scan(execution) - + python_factory = partial(_python_vector_scan, execution) if not _native_requested(builder): return python_factory() from pypaimon.execution.native.read.native_search_plan import native_vector_scan @@ -90,10 +109,7 @@ def python_factory(): def create_vector_read(builder, batch=False): - def python_factory(): - from pypaimon.execution.python.read.search.factories import new_batch_vector_read, new_vector_read - return new_batch_vector_read(builder) if batch else new_vector_read(builder) - + python_factory = partial(_python_vector_read, builder, batch) if not _native_requested(builder): return python_factory() from pypaimon.execution.native.read.native_search_plan import native_vector_read @@ -105,10 +121,7 @@ def python_factory(): def create_full_text_scan(builder, execution): - def python_factory(): - from pypaimon.execution.python.read.search.factories import new_full_text_scan - return new_full_text_scan(execution) - + python_factory = partial(_python_full_text_scan, execution) if not _native_requested(builder): return python_factory() from pypaimon.execution.native.read.native_search_plan import native_full_text_scan @@ -117,10 +130,7 @@ def python_factory(): def create_full_text_read(builder, execution): - def python_factory(): - from pypaimon.execution.python.read.search.factories import new_full_text_read - return new_full_text_read(execution) - + python_factory = partial(_python_full_text_read, execution) if not _native_requested(builder): return python_factory() from pypaimon.execution.native.read.native_search_plan import native_full_text_read diff --git a/paimon-python/pypaimon/tests/multimodal_lance_descriptor_test.py b/paimon-python/pypaimon/tests/multimodal_lance_descriptor_test.py index 46c508406e64..5ea094841997 100644 --- a/paimon-python/pypaimon/tests/multimodal_lance_descriptor_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lance_descriptor_test.py @@ -119,9 +119,10 @@ def test_export_lance_rejects_video_frame_in_ordinary_descriptor_field(conn, tmp video = tmp_path / "video.bin" video.write_bytes(b"encoded-video") descriptor = VideoFrameDescriptor(str(video), 0, video.stat().st_size, 2, -1, 0) + # Seed the invalid legacy value through Python; Native rejects it before export. table = lance_tests._table(conn, "frame_reference", {"id": pa.int32(), "payload": pa.large_binary()}, [ {"id": 1, "payload": descriptor}, - ], options={"blob-descriptor-field": "payload"}) + ], options={"blob-descriptor-field": "payload", "write.native.enabled": "false"}) path = tmp_path / "frame_reference.lance" if mode == "overwrite": lance.write_dataset(pa.table({"id": [99]}), str(path)) diff --git a/paimon-python/pypaimon/tests/native_blob_update_test.py b/paimon-python/pypaimon/tests/native_blob_update_test.py index 19e2968f2b6f..b434fe50155e 100644 --- a/paimon-python/pypaimon/tests/native_blob_update_test.py +++ b/paimon-python/pypaimon/tests/native_blob_update_test.py @@ -18,7 +18,9 @@ """Native updates share the Java Blob placeholder and row-id layout.""" from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId -from contextlib import ExitStack +from pypaimon.execution.python.write.table_update import PythonTableUpdate +from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert +from contextlib import ExitStack, contextmanager from unittest.mock import patch from urllib.parse import urlparse @@ -31,11 +33,8 @@ from pypaimon.table.row.blob import Blob, BlobDescriptor, BlobViewStruct from pypaimon.table.row.generic_row import GenericRow from pypaimon.table.data_evolution_merge_into import WhenMatched, WhenNotMatched, source_col -from pypaimon.write.table_delete import TableDeleteByRowId from pypaimon.execution.native.write.native_update import NativeTableUpdateByRowId from pypaimon.write.write_dispatch import NativeTableWrite -from pypaimon.write.table_update import TableUpdate -from pypaimon.write.table_upsert_by_key import TableUpsertByKey pytestmark = pytest.mark.native_plan @@ -85,15 +84,17 @@ def _read(table, descriptors=False, projection=None, resolve_views=True): return read.to_arrow(splits).sort_by('id').to_pylist() +@contextmanager def _native_only(): - stack = ExitStack() - for cls, method in [(PythonTableUpdateByRowId, '_load_existing_files_info'), - (TableUpdate, '_build_predicate_update_table'), - (TableUpsertByKey, '_upsert_partition'), - (TableUpsertByKey, '_upsert_row_partition'), - (TableDeleteByRowId, 'delete')]: - stack.enter_context(patch.object(cls, method, side_effect=AssertionError('Python update fallback'))) - return stack + with ExitStack() as stack: + for cls, method in [(PythonTableUpdateByRowId, '_load_existing_files_info'), + (PythonTableUpdate, '_build_predicate_update_table'), + (PythonTableUpsert, '_upsert_partition'), + (PythonTableUpsert, '_upsert_row_partition'), + (PythonTableUpdate, 'delete'), + (PythonTableUpdate, 'delete_by_row_id')]: + stack.enter_context(patch.object(cls, method, side_effect=AssertionError('Python update fallback'))) + yield def _commit(builder, messages, stream): @@ -238,13 +239,13 @@ def test_native_inline_blob_updates_preserve_references( else update.update_by_arrow_with_row_id(data)) elif operation == 'incremental': updater = update.new_update_by_row_id(7) if stream else update.new_update_by_row_id() - assert isinstance(updater, NativeTableUpdateByRowId) == native + assert isinstance(updater._execution, NativeTableUpdateByRowId) == native updater.update_columns(data.slice(0, 1), ['payload']) updater.update_columns(data.slice(1), ['payload']) messages = updater.commit_messages elif operation in ('row', 'rows'): updater = update.new_update_by_row_id(7) if stream else update.new_update_by_row_id() - assert isinstance(updater, NativeTableUpdateByRowId) == native + assert isinstance(updater._execution, NativeTableUpdateByRowId) == native value = (Blob.from_view(BlobViewStruct.deserialize(replacement)) if option == 'blob-view-field' else Blob.from_descriptor_bytes(replacement, file_io=table.file_io)) rows = [GenericRow([1, 'a', 999, value], table.fields), @@ -347,13 +348,13 @@ def test_native_raw_blob_updates_across_existing_update_apis(tmp_path, kind, str messages = update.update_by_arrow_batches_with_row_id(batches) elif operation == 'incremental': updater = update.new_update_by_row_id(7) if stream else update.new_update_by_row_id() - assert isinstance(updater, NativeTableUpdateByRowId) + assert isinstance(updater._execution, NativeTableUpdateByRowId) updater.update_columns(data.slice(0, 1), ['payload']) updater.update_columns(data.slice(1), ['payload']) messages = updater.commit_messages elif operation in ('row', 'rows'): updater = update.new_update_by_row_id(7) if stream else update.new_update_by_row_id() - assert isinstance(updater, NativeTableUpdateByRowId) + assert isinstance(updater._execution, NativeTableUpdateByRowId) rows = [GenericRow([1, 'a', 999, replacement], table.fields), GenericRow([3, 'b', 999, None], table.fields)] if operation == 'row': @@ -499,8 +500,8 @@ def test_raw_blob_abort_preserves_prepared_files(tmp_path, published): paths = set(tmp_path.rglob('*.blob')) if published: _commit(builder, messages, False) - updater.writer._abort() - updater.writer._abort() + updater._execution.writer._abort() + updater._execution.writer._abort() assert set(tmp_path.rglob('*.blob')) == paths if not published: assert _read(table)[0]['payload'] == b'large' * 10 diff --git a/paimon-python/pypaimon/tests/native_map_shredding_write_test.py b/paimon-python/pypaimon/tests/native_map_shredding_write_test.py index df647e636176..52fe2f5ea27c 100644 --- a/paimon-python/pypaimon/tests/native_map_shredding_write_test.py +++ b/paimon-python/pypaimon/tests/native_map_shredding_write_test.py @@ -278,8 +278,8 @@ def test_map_row_id_update_keeps_physical_layout(tmp_path, native_update): def test_native_map_predicate_update_and_upsert(tmp_path, stream, operation): from contextlib import ExitStack from pypaimon.common.predicate_builder import PredicateBuilder - from pypaimon.write.table_update import TableUpdate - from pypaimon.write.table_upsert_by_key import TableUpsertByKey + from pypaimon.execution.python.write.table_update import PythonTableUpdate + from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert table, schema = _table(tmp_path, 'evolution') _write(table, schema, [ @@ -291,8 +291,8 @@ def test_native_map_predicate_update_and_upsert(tmp_path, stream, operation): replacement = [('new', 7), ('a', None)] with ExitStack() as stack: for cls, method in [(PythonTableUpdateByRowId, '_load_existing_files_info'), - (TableUpdate, '_build_predicate_update_table'), - (TableUpsertByKey, '_upsert_partition')]: + (PythonTableUpdate, '_build_predicate_update_table'), + (PythonTableUpsert, '_upsert_partition')]: stack.enter_context(patch.object(cls, method, side_effect=AssertionError('Python update fallback'))) if operation == 'upsert': data = pa.Table.from_pylist([ diff --git a/paimon-python/pypaimon/tests/native_nested_update_test.py b/paimon-python/pypaimon/tests/native_nested_update_test.py index ed5df8a02df0..1eea346dc090 100644 --- a/paimon-python/pypaimon/tests/native_nested_update_test.py +++ b/paimon-python/pypaimon/tests/native_nested_update_test.py @@ -18,7 +18,7 @@ """Nested Native update dispatch and parity with the Python implementation.""" from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId -from contextlib import ExitStack +from contextlib import ExitStack, contextmanager from unittest.mock import patch import pyarrow as pa @@ -27,8 +27,8 @@ from pypaimon import CatalogFactory, Schema from pypaimon.common.predicate_builder import PredicateBuilder from pypaimon.table.row.generic_row import GenericRow -from pypaimon.write.table_update import TableUpdate -from pypaimon.write.table_upsert_by_key import TableUpsertByKey +from pypaimon.execution.python.write.table_update import PythonTableUpdate +from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert pytestmark = pytest.mark.native_plan @@ -66,15 +66,16 @@ def _table(tmp_path, data_type, native): return table, schema +@contextmanager def _native_only(native): - stack = ExitStack() - if native: - for cls, method in [(PythonTableUpdateByRowId, '_load_existing_files_info'), - (TableUpdate, '_build_predicate_update_table'), - (TableUpsertByKey, '_upsert_partition'), - (TableUpsertByKey, '_upsert_row_partition')]: - stack.enter_context(patch.object(cls, method, side_effect=AssertionError('Python fallback'))) - return stack + with ExitStack() as stack: + if native: + for cls, method in [(PythonTableUpdateByRowId, '_load_existing_files_info'), + (PythonTableUpdate, '_build_predicate_update_table'), + (PythonTableUpsert, '_upsert_partition'), + (PythonTableUpsert, '_upsert_row_partition')]: + stack.enter_context(patch.object(cls, method, side_effect=AssertionError('Python fallback'))) + yield def _commit(builder, messages, stream): diff --git a/paimon-python/pypaimon/tests/native_parquet_vector_update_test.py b/paimon-python/pypaimon/tests/native_parquet_vector_update_test.py index 06f3c8dc7d1e..43f1bf42272a 100644 --- a/paimon-python/pypaimon/tests/native_parquet_vector_update_test.py +++ b/paimon-python/pypaimon/tests/native_parquet_vector_update_test.py @@ -18,7 +18,7 @@ """Parquet vector sidecars must not force row-level updates back to Python.""" from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId -from contextlib import ExitStack +from contextlib import ExitStack, contextmanager from unittest.mock import patch import pyarrow as pa @@ -29,8 +29,8 @@ from pypaimon.execution.native.read.table_read import (native_read as read_native_splits) from pypaimon.schema.data_types import PyarrowFieldParser from pypaimon.table.data_evolution_merge_into import WhenMatched, WhenNotMatched, source_col -from pypaimon.write.table_update import TableUpdate -from pypaimon.write.table_upsert_by_key import TableUpsertByKey +from pypaimon.execution.python.write.table_update import PythonTableUpdate +from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert pytestmark = pytest.mark.native_plan @@ -73,16 +73,17 @@ def _rows(table): yield sorted(rows, key=lambda row: row['id']) +@contextmanager def _native_only(native): - stack = ExitStack() - if native: - for owner, method in [(PythonTableUpdateByRowId, '_load_existing_files_info'), - (TableUpsertByKey, '_upsert_partition'), - (TableUpdate, '_build_predicate_update_table')]: - stack.enter_context(patch.object(owner, method, side_effect=AssertionError('Python update fallback'))) - stack.enter_context(patch('pypaimon.execution.python.write.merge_into._build_tables', - side_effect=AssertionError('Python merge fallback'))) - return stack + with ExitStack() as stack: + if native: + for owner, method in [(PythonTableUpdateByRowId, '_load_existing_files_info'), + (PythonTableUpsert, '_upsert_partition'), + (PythonTableUpdate, '_build_predicate_update_table')]: + stack.enter_context(patch.object(owner, method, side_effect=AssertionError('Python update fallback'))) + stack.enter_context(patch('pypaimon.execution.python.write.merge_into._build_tables', + side_effect=AssertionError('Python merge fallback'))) + yield def _commit(builder, messages, stream): diff --git a/paimon-python/pypaimon/tests/native_ray_index_build_test.py b/paimon-python/pypaimon/tests/native_ray_index_build_test.py index 5f6fd1b6261d..d7de6107812a 100644 --- a/paimon-python/pypaimon/tests/native_ray_index_build_test.py +++ b/paimon-python/pypaimon/tests/native_ray_index_build_test.py @@ -96,7 +96,7 @@ def _request(tmp_path): @pytest.fixture -def native_protocol(): +def native_protocol(monkeypatch): module = 'pypaimon.execution.native.write.native_index_build' core = Mock() shards = [SimpleNamespace(ordinal=ordinal) for ordinal in range(3)] @@ -106,8 +106,11 @@ def native_protocol(): sdk = ModuleType('pypaimon_rust.datafusion') sdk.GlobalIndexBuildBuilder = SimpleNamespace(supports_index_type=lambda kind: True) scheduler = _Ray() - with patch.dict(sys.modules, {'pypaimon_rust.datafusion': sdk, 'ray': scheduler}), \ - patch(module + '.native_runtime_available', return_value=True), \ + # Retain lazy imports between cases instead of restoring all of sys.modules, + # which leaves stale package attributes for Python 3.10's mock resolver. + monkeypatch.setitem(sys.modules, 'pypaimon_rust.datafusion', sdk) + monkeypatch.setitem(sys.modules, 'ray', scheduler) + with patch(module + '.native_runtime_available', return_value=True), \ patch(module + '._rest_catalog_supported', return_value=True), \ patch(module + '._native_table', return_value=native_table), \ patch(module + '.from_native_commit_messages', side_effect=lambda table, messages: messages), \ diff --git a/paimon-python/pypaimon/tests/native_row_sidecar_test.py b/paimon-python/pypaimon/tests/native_row_sidecar_test.py index 7e91ccea0d07..5f514b9cf8b3 100644 --- a/paimon-python/pypaimon/tests/native_row_sidecar_test.py +++ b/paimon-python/pypaimon/tests/native_row_sidecar_test.py @@ -203,7 +203,7 @@ def test_sidecars_keep_updates_native_and_apply_deletion_vectors(tmp_path, strea data = pa.table({'_ROW_ID': [5, 7], 'value': ['changed-5', 'changed-7']}) if incremental: updater = update.new_update_by_row_id(1) if stream else update.new_update_by_row_id() - assert isinstance(updater, NativeTableUpdateByRowId) + assert isinstance(updater._execution, NativeTableUpdateByRowId) messages = updater.update_columns(data, ['value']) else: with patch.object(TableUpdateByRowId, 'update_columns', diff --git a/paimon-python/pypaimon/tests/native_update_container_casts_test.py b/paimon-python/pypaimon/tests/native_update_container_casts_test.py index 3b2be5e9842d..fc043a9cf20a 100644 --- a/paimon-python/pypaimon/tests/native_update_container_casts_test.py +++ b/paimon-python/pypaimon/tests/native_update_container_casts_test.py @@ -18,7 +18,7 @@ """Equivalent Arrow container inputs must work on the existing update APIs.""" from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId -from contextlib import ExitStack +from contextlib import ExitStack, contextmanager from decimal import Decimal from unittest.mock import patch @@ -28,8 +28,8 @@ from pypaimon import CatalogFactory, Schema from pypaimon.common.predicate_builder import PredicateBuilder from pypaimon.table.data_evolution_merge_into import WhenMatched, source_col -from pypaimon.write.table_update import TableUpdate -from pypaimon.write.table_upsert_by_key import TableUpsertByKey +from pypaimon.execution.python.write.table_update import PythonTableUpdate +from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert pytestmark = pytest.mark.native_plan @@ -92,16 +92,17 @@ def _table(tmp_path, value_type, native): return table +@contextmanager def _native_only(native): - stack = ExitStack() - if native: - for cls, method in [(PythonTableUpdateByRowId, '_load_existing_files_info'), - (TableUpdate, '_build_predicate_update_table'), - (TableUpsertByKey, '_upsert_partition')]: - stack.enter_context(patch.object(cls, method, side_effect=AssertionError('Python update fallback'))) - stack.enter_context(patch('pypaimon.execution.python.write.merge_into._build_tables', - side_effect=AssertionError('Python MERGE fallback'))) - return stack + with ExitStack() as stack: + if native: + for cls, method in [(PythonTableUpdateByRowId, '_load_existing_files_info'), + (PythonTableUpdate, '_build_predicate_update_table'), + (PythonTableUpsert, '_upsert_partition')]: + stack.enter_context(patch.object(cls, method, side_effect=AssertionError('Python update fallback'))) + stack.enter_context(patch('pypaimon.execution.python.write.merge_into._build_tables', + side_effect=AssertionError('Python MERGE fallback'))) + yield def _apply(table, source, native, stream, operation): diff --git a/paimon-python/pypaimon/tests/native_update_inactive_vector_option_test.py b/paimon-python/pypaimon/tests/native_update_inactive_vector_option_test.py index 083845654f70..e12920c8295e 100644 --- a/paimon-python/pypaimon/tests/native_update_inactive_vector_option_test.py +++ b/paimon-python/pypaimon/tests/native_update_inactive_vector_option_test.py @@ -19,7 +19,7 @@ from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId from pypaimon.execution.python.write.table_delete import PythonTableDeleteByRowId -from contextlib import ExitStack +from contextlib import ExitStack, contextmanager from unittest.mock import patch import pyarrow as pa @@ -32,8 +32,8 @@ from pypaimon.table.row.generic_row import GenericRow from pypaimon.tests import native_plan_rest_test from pypaimon.execution.native.write.native_update import _native_update_formats_supported -from pypaimon.write.table_update import TableUpdate -from pypaimon.write.table_upsert_by_key import TableUpsertByKey +from pypaimon.execution.python.write.table_update import PythonTableUpdate +from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert rest_catalog = native_plan_rest_test.rest_catalog @@ -80,18 +80,19 @@ def _read(table): return builder.new_read().to_arrow(builder.new_scan().plan().splits()).sort_by('id').to_pylist() +@contextmanager def _native_only(native): - stack = ExitStack() - if native: - for owner, method in [(PythonTableUpdateByRowId, '_load_existing_files_info'), - (TableUpsertByKey, '_upsert_partition'), - (TableUpsertByKey, '_upsert_row_partition'), - (TableUpdate, '_build_predicate_update_table'), - (PythonTableDeleteByRowId, '_scan_anchor_ranges')]: - stack.enter_context(patch.object(owner, method, side_effect=AssertionError('Python mutation fallback'))) - stack.enter_context(patch('pypaimon.execution.python.write.merge_into._build_tables', - side_effect=AssertionError('Python MERGE fallback'))) - return stack + with ExitStack() as stack: + if native: + for owner, method in [(PythonTableUpdateByRowId, '_load_existing_files_info'), + (PythonTableUpsert, '_upsert_partition'), + (PythonTableUpsert, '_upsert_row_partition'), + (PythonTableUpdate, '_build_predicate_update_table'), + (PythonTableDeleteByRowId, '_scan_anchor_ranges')]: + stack.enter_context(patch.object(owner, method, side_effect=AssertionError('Python mutation fallback'))) + stack.enter_context(patch('pypaimon.execution.python.write.merge_into._build_tables', + side_effect=AssertionError('Python MERGE fallback'))) + yield @pytest.mark.parametrize('native', _ENGINES, ids=['python', 'native']) diff --git a/paimon-python/pypaimon/tests/native_upsert_layout_test.py b/paimon-python/pypaimon/tests/native_upsert_layout_test.py index 223fda60793c..189080d8a991 100644 --- a/paimon-python/pypaimon/tests/native_upsert_layout_test.py +++ b/paimon-python/pypaimon/tests/native_upsert_layout_test.py @@ -18,7 +18,7 @@ """Equivalent Arrow layouts must not send upsert matching back to Python.""" from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId -from contextlib import ExitStack +from contextlib import ExitStack, contextmanager from unittest.mock import patch import pyarrow as pa @@ -26,8 +26,8 @@ from pypaimon import CatalogFactory, Schema from pypaimon.schema.data_types import PyarrowFieldParser -from pypaimon.write.table_update import TableUpdate -from pypaimon.write.table_upsert_by_key import TableUpsertByKey +from pypaimon.execution.python.write.table_update import PythonTableUpdate +from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert pytestmark = pytest.mark.native_plan @@ -76,14 +76,15 @@ def _seed(table, schema, keys, values): writer.close() +@contextmanager def _native_only(native): - stack = ExitStack() - if native: - for owner, method in [(TableUpsertByKey, '_upsert_partition'), - (PythonTableUpdateByRowId, '_load_existing_files_info'), - (TableUpdate, '_build_predicate_update_table')]: - stack.enter_context(patch.object(owner, method, side_effect=AssertionError('Python update fallback'))) - return stack + with ExitStack() as stack: + if native: + for owner, method in [(PythonTableUpsert, '_upsert_partition'), + (PythonTableUpdateByRowId, '_load_existing_files_info'), + (PythonTableUpdate, '_build_predicate_update_table')]: + stack.enter_context(patch.object(owner, method, side_effect=AssertionError('Python update fallback'))) + yield def _check(table, values, key, new_key='new'): diff --git a/paimon-python/pypaimon/tests/native_variant_update_test.py b/paimon-python/pypaimon/tests/native_variant_update_test.py index 0335cf9956db..6f0f5497c535 100644 --- a/paimon-python/pypaimon/tests/native_variant_update_test.py +++ b/paimon-python/pypaimon/tests/native_variant_update_test.py @@ -19,7 +19,7 @@ from pypaimon.execution.python.write.table_update_by_row_id import PythonTableUpdateByRowId import json -from contextlib import ExitStack +from contextlib import ExitStack, contextmanager from unittest.mock import patch import pyarrow as pa @@ -32,8 +32,8 @@ from pypaimon.schema.data_types import PyarrowFieldParser from pypaimon.execution.native.write.native_update import NativeTableUpdateByRowId from pypaimon.write.write_dispatch import NativeTableWrite -from pypaimon.write.table_update import TableUpdate -from pypaimon.write.table_upsert_by_key import TableUpsertByKey +from pypaimon.execution.python.write.table_update import PythonTableUpdate +from pypaimon.execution.python.write.table_upsert_by_key import PythonTableUpsert pytestmark = pytest.mark.native_plan @@ -104,13 +104,14 @@ def _read(table, native, projection=None): return rows +@contextmanager def _native_only(): - stack = ExitStack() - for cls, method in [(PythonTableUpdateByRowId, '_load_existing_files_info'), - (TableUpdate, '_build_predicate_update_table'), - (TableUpsertByKey, '_upsert_partition')]: - stack.enter_context(patch.object(cls, method, side_effect=AssertionError('Python update fallback'))) - return stack + with ExitStack() as stack: + for cls, method in [(PythonTableUpdateByRowId, '_load_existing_files_info'), + (PythonTableUpdate, '_build_predicate_update_table'), + (PythonTableUpsert, '_upsert_partition')]: + stack.enter_context(patch.object(cls, method, side_effect=AssertionError('Python update fallback'))) + yield def test_python_configured_variant_schema_uses_java_semantics(tmp_path): @@ -145,7 +146,7 @@ def test_native_variant_partial_updates(tmp_path, stream, key, enabled, operatio else update.update_by_arrow_with_row_id(data)) elif operation == 'incremental': updater = update.new_update_by_row_id(7) if stream else update.new_update_by_row_id() - assert isinstance(updater, NativeTableUpdateByRowId) + assert isinstance(updater._execution, NativeTableUpdateByRowId) updater.update_columns(data.slice(0, 1), ['payload']) updater.update_columns(data.slice(1), ['payload']) messages = updater.commit_messages 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 04fd3e5d9803..5392d27ad0dc 100644 --- a/paimon-python/pypaimon/tests/native_vector_index_build_test.py +++ b/paimon-python/pypaimon/tests/native_vector_index_build_test.py @@ -137,8 +137,7 @@ def test_native_search_filters_live_null_vectors_with_deletion_vectors(rest_cata result = (_native_table(table).new_vector_search_builder().with_vector_column('embedding') .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. + # Python table-level reads use the batch ABI for short filtered results. _assert_ids(table, [1, 4], kind) diff --git a/paimon-python/pypaimon/tests/native_write_test.py b/paimon-python/pypaimon/tests/native_write_test.py index 48d3b7765c2d..ecc2f1420fc5 100644 --- a/paimon-python/pypaimon/tests/native_write_test.py +++ b/paimon-python/pypaimon/tests/native_write_test.py @@ -528,7 +528,7 @@ def test_uri_reader_factory_none_keeps_native_and_rejects_late_change(tmp_path): native.write_arrow_batch(_batch([2], ['a'])) with pytest.raises(RuntimeError, match='before any write'): native.with_blob_uri_reader_factory(None) - with pytest.raises(ValueError, match='before any data is written'): + with pytest.raises(RuntimeError, match='before any write'): native.with_write_type(['id']) native.abort() diff --git a/paimon-python/pypaimon/tests/raw_file_scan_auth_scope_test.py b/paimon-python/pypaimon/tests/raw_file_scan_auth_scope_test.py index d1cbe980f8f4..83566f6e6763 100644 --- a/paimon-python/pypaimon/tests/raw_file_scan_auth_scope_test.py +++ b/paimon-python/pypaimon/tests/raw_file_scan_auth_scope_test.py @@ -92,14 +92,22 @@ def authorize(select): assert native.call_count == int(native_fallback) -def test_restricted_raw_scope_does_not_clear_cached_rules(table): +@pytest.mark.parametrize('operation', ['raw-files', 'all-files']) +def test_restricted_raw_scope_does_not_clear_cached_rules(table, operation): scan = table.new_read_builder().new_scan() scan._query_auth_fn = lambda select: TableQueryAuthResult([_partition_filter()], None) scan.plan() source = scan._python_file_scanner prior_partition_filter = source.auth_partition_predicate assert prior_partition_filter is not None - with patch.object(source, 'plan_files', side_effect=AssertionError('Unauthorized file access')): + plan = scan.file_scanner.plan_files if operation == 'raw-files' else scan.plan_all + with patch.object(source, 'plan_files', side_effect=AssertionError('Unauthorized file access')), \ + patch.object(NativeTableScan, 'plan_files', + side_effect=AssertionError('Unauthorized Native file access')) as native, \ + patch.object(scan, '_try_native_plan', + side_effect=AssertionError('Unauthorized Native planning')) as native_plan: with pytest.raises(TableNoPermissionException): - scan.file_scanner.plan_files() + plan() + native.assert_not_called() + native_plan.assert_not_called() assert source.auth_partition_predicate is prior_partition_filter diff --git a/paimon-python/pypaimon/tests/ray_bucket_join_test.py b/paimon-python/pypaimon/tests/ray_bucket_join_test.py index 292712361d13..115be8aae553 100644 --- a/paimon-python/pypaimon/tests/ray_bucket_join_test.py +++ b/paimon-python/pypaimon/tests/ray_bucket_join_test.py @@ -309,7 +309,7 @@ def test_rejects_rescaled_mixed_buckets(self): # Files left under an old bucket count (rescale not yet rewritten) would carry a # different total_buckets; the same bucket id must not be treated as co-located. import types - from pypaimon.execution.python.read.scanner.file_scanner import (FileScanner) + from pypaimon.read.file_scanner import FileScanner sch = pa.schema([("url", pa.string())]) self._bucketed_table("default.rs_a", sch, "url", pa.Table.from_pydict({"url": ["u0", "u1"]}, schema=sch)) diff --git a/paimon-python/pypaimon/tests/read_execution_boundary_test.py b/paimon-python/pypaimon/tests/read_execution_boundary_test.py index 4f1236e7bcb2..878a2b05f877 100644 --- a/paimon-python/pypaimon/tests/read_execution_boundary_test.py +++ b/paimon-python/pypaimon/tests/read_execution_boundary_test.py @@ -17,6 +17,7 @@ """Read entry points load only the selected execution backend.""" +import importlib import subprocess import sys import textwrap @@ -259,7 +260,8 @@ def test_native_row_conversion_keeps_its_failure_when_closing_batches_also_fails read.output_projection = None read._output_arrow_schema = lambda splits: pa.schema([('id', pa.int32())]) read._try_native_batches = Mock(return_value=source) - with patch('pypaimon.execution.native.read.row_iterator.native_rows', side_effect=primary), \ + row_iterator = importlib.import_module('pypaimon.execution.native.read.row_iterator') + with patch.object(row_iterator, 'native_rows', side_effect=primary), \ patch.dict(sys.modules, {'pypaimon.execution.python.read.table_read': None}): rows = read._try_native_rows([]) with pytest.raises(ValueError) as caught: @@ -272,6 +274,7 @@ def _python_read_source(primary=None, cleanup=None): from types import SimpleNamespace import pyarrow as pa from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader + from pypaimon.read.read_support import _try_to_pad_batch_by_schema from pypaimon.schema.data_types import AtomicType, DataField from pypaimon.table.row.offset_row import OffsetRow @@ -310,6 +313,7 @@ def close(self): _blob_as_descriptor=False, _blob_view_output_indices=(), _project_batch_to_output=lambda batch: batch, + _try_to_pad_batch_by_schema=_try_to_pad_batch_by_schema, output_projection=None, _adapter_read_type=[DataField(0, 'id', AtomicType('INT'))], include_row_kind=False, diff --git a/paimon-python/pypaimon/tests/vindex_vector_index_test.py b/paimon-python/pypaimon/tests/vindex_vector_index_test.py index c2d67bec7da1..46a83ae2839d 100644 --- a/paimon-python/pypaimon/tests/vindex_vector_index_test.py +++ b/paimon-python/pypaimon/tests/vindex_vector_index_test.py @@ -19,17 +19,21 @@ import sys import types import unittest -from unittest.mock import patch +from unittest.mock import Mock, patch + +import numpy as np from pypaimon.globalindex.batch_vector_search import BatchVectorSearch from pypaimon.globalindex.global_index_meta import GlobalIndexIOMeta from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult +from pypaimon.globalindex.vector_search import VectorSearch from pypaimon.globalindex.vindex_reader import ( VindexVectorGlobalIndexReader, ) from pypaimon.execution.python.globalindex.vindex.vindex_vector_global_index_reader import ( _build_scores, _search_params, + PythonVindexVectorGlobalIndexReader, ) from pypaimon.utils.roaring_bitmap import RoaringBitmap64 @@ -113,6 +117,23 @@ def new_input_stream(self, path): class VindexVectorIndexTest(unittest.TestCase): + def test_single_query_preserves_short_filtered_results(self): + reader = PythonVindexVectorGlobalIndexReader( + _BytesFileIO(), '/tmp', [GlobalIndexIOMeta('index', 1)]) + reader._metadata = _FakeMetadata() + reader._search_params_type = _FakeSearchParams + reader._reader = Mock() + reader._reader.search.side_effect = RuntimeError('scalar ABI requires exactly K results') + reader._reader.search_batch.return_value = ( + np.array([[10, 11, -1]]), np.array([[0., 1., np.nan]])) + with patch.object(reader, '_ensure_loaded'): + result = reader.visit_vector_search(VectorSearch([1., 0.], 3, 'embedding')).result() + self.assertEqual([10, 11], result.results().to_list()) + self.assertEqual(1., result.score_getter()(10)) + self.assertEqual(.5, result.score_getter()(11)) + reader._reader.search.assert_not_called() + reader.close() + def test_query_widths_match_java_integer_text_and_bounds(self): for key in ('ivf.nprobe', 'diskann.l_search'): for value, expected in (('+12', 12), ('12', 12), ('٢', 2), ('2147483647', 2147483647), (12, 12)): diff --git a/paimon-python/pypaimon/tests/write/execution_lifecycle_test.py b/paimon-python/pypaimon/tests/write/execution_lifecycle_test.py index 1c3938cfb580..3272eab1749a 100644 --- a/paimon-python/pypaimon/tests/write/execution_lifecycle_test.py +++ b/paimon-python/pypaimon/tests/write/execution_lifecycle_test.py @@ -267,7 +267,10 @@ def test_native_merge_does_not_construct_python_writer_or_file_index(tmp_path): table = _data_evolution_table(tmp_path) native_table, builder = _core_builder() expected = [_message()] + encoded_clause = object() with patch('pypaimon.execution.native.write.native_merge_into._native_row_id_table', return_value=native_table), \ + patch('pypaimon.execution.native.write.native_merge_into.encode_merge_clause', + return_value=encoded_clause), \ patch('pypaimon.execution.native.write.native_merge_into.from_native_commit_messages', return_value=expected), \ patch('pypaimon.execution.python.write.table_write.PythonTableWrite.__init__', @@ -279,6 +282,7 @@ def test_native_merge_does_not_construct_python_writer_or_file_index(tmp_path): when_matched=[WhenMatched.update({'value': source_col('value')})]) assert result is expected builder.new_update.return_value.merge_into.assert_called_once() + assert builder.new_update.return_value.merge_into.call_args[1]['when_matched'] == [encoded_clause] def test_failed_python_construction_closes_snapshot_and_preserves_original_error(): diff --git a/paimon-python/pypaimon/tests/write/row_id_conflict_protocol_test.py b/paimon-python/pypaimon/tests/write/row_id_conflict_protocol_test.py index 091cc57d2057..4c22822587eb 100644 --- a/paimon-python/pypaimon/tests/write/row_id_conflict_protocol_test.py +++ b/paimon-python/pypaimon/tests/write/row_id_conflict_protocol_test.py @@ -91,7 +91,8 @@ def test_native_commit_translates_only_typed_row_id_conflicts(batch): @pytest.mark.parametrize('opaque_message', ['unknown commit result', 'Row ID existence conflict: not a typed error']) -def test_native_unknown_outcome_is_not_a_rewrite_signal(opaque_message): +@pytest.mark.parametrize('native_available', [False, True]) +def test_native_unknown_outcome_is_not_a_rewrite_signal(opaque_message, native_available): from pypaimon.execution.native.write.commit import NativeTableCommit from pypaimon.ray.row_id_conflict_rewriter import _find_row_id_conflict @@ -100,7 +101,7 @@ def test_native_unknown_outcome_is_not_a_rewrite_signal(opaque_message): failure.__cause__ = _native_conflict(module) core = Mock() core.commit.side_effect = failure - with patch.dict(sys.modules, {'pypaimon_rust': module}): + with patch.dict(sys.modules, {'pypaimon_rust': module if native_available else None}): with pytest.raises(ValueError) as raised: NativeTableCommit(core, is_batch=True).commit([], 7) assert raised.value is failure