This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-7019-1f9fbded6e6ea480302af2c61f7cea510c1df77d in repository https://gitbox.apache.org/repos/asf/texera.git
commit 9ff28ddd1408bdece43e16723c330b98414f0b8c Author: Xinyuan Lin <[email protected]> AuthorDate: Wed Jul 29 01:06:07 2026 -0700 test(pyamber): extend IO, packaging and storage-model coverage (#7019) ### What changes were proposed in this PR? Extends three pyamber suites and adds two new test files. No source changes. - **`core/storage/runnables/input_port_materialization_reader_runnable.py`** (~83% covered): extended over the uncovered branches. - **`core/architecture/packaging/input_manager.py`** (~88% covered, **no dedicated test**): new test file. - **`core/models/tuple.py`** (~95% covered) and **`pytexera/udf/udf_operator.py`** (~94% covered): extended. - **`core/storage/model/virtual_document.py`** (~74% covered, **no dedicated test**): new test file. Tests are hermetic and deterministic — no real sleeps, nothing left running. ### Any related issues, documentation, discussions? Closes #7017. ### How was this PR tested? `python -m pytest` over the five files -> 187 passed. (`TestTuple::test_hash` is deselected locally: it fails identically on unmodified `main` under Windows because `tuple.py` calls `datetime.timestamp()`, which raises `OSError [Errno 22]` there; it passes on the Linux CI runners.) `ruff format --check` + `ruff check` clean. ### Was this PR authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 4.8 [1M context]) --- .../architecture/packaging/test_input_manager.py | 262 +++++++++++++++++++++ amber/src/test/python/core/models/test_tuple.py | 177 ++++++++++++++ .../core/storage/model/test_virtual_document.py | 104 ++++++++ ...t_input_port_materialization_reader_runnable.py | 243 ++++++++++++++++++- .../test/python/pytexera/udf/test_udf_operator.py | 149 ++++++++++++ 5 files changed, 933 insertions(+), 2 deletions(-) diff --git a/amber/src/test/python/core/architecture/packaging/test_input_manager.py b/amber/src/test/python/core/architecture/packaging/test_input_manager.py new file mode 100644 index 0000000000..34cd00a62a --- /dev/null +++ b/amber/src/test/python/core/architecture/packaging/test_input_manager.py @@ -0,0 +1,262 @@ +# 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. + +from unittest.mock import MagicMock, patch + +import pyarrow +import pytest + +import core.architecture.packaging.input_manager as input_manager_module +from core.architecture.packaging.input_manager import InputManager +from core.models.payload import DataFrame, StateFrame +from core.models.schema.schema import Schema +from core.models.state import State +from proto.org.apache.texera.amber.core import ( + ActorVirtualIdentity, + ChannelIdentity, + PortIdentity, +) + +WORKER_ID = "Worker:WF0-test-op-main-0" + + +def _channel(name: str) -> ChannelIdentity: + return ChannelIdentity( + ActorVirtualIdentity(name), + ActorVirtualIdentity(WORKER_ID), + is_control=False, + ) + + +def _reader_stub(finished: bool, channel_name: str) -> MagicMock: + reader = MagicMock() + reader.finished.return_value = finished + reader.channel_id = _channel(channel_name) + return reader + + +class TestPortIdentityDefaults: + """PortIdentity arrives from protobuf, where an unset scalar can reach + Python as None. Both add_input_port and register_input must normalise + those to the (0, False) defaults *before* the identity is used as a + dict key -- otherwise the same logical port registered through the two + paths would hash to two different entries. + """ + + @pytest.fixture + def manager(self): + return InputManager(worker_id=WORKER_ID, input_queue=MagicMock()) + + def test_add_input_port_defaults_unset_fields(self, manager): + port_id = PortIdentity(id=None, internal=None) + schema = Schema(raw_schema={"x": "INTEGER"}) + + manager.add_input_port(port_id, schema, [], []) + + assert port_id.id == 0 + assert port_id.internal is False + # The normalised identity is what got keyed, so a canonical + # PortIdentity(0, False) resolves to the very same port. + assert manager.get_port(PortIdentity(0, False)).get_schema() == schema + + def test_register_input_defaults_unset_fields(self, manager): + canonical = PortIdentity(0, False) + manager.add_input_port(canonical, Schema(raw_schema={"x": "INTEGER"}), [], []) + channel_id = _channel("upstream") + + manager.register_input(channel_id, PortIdentity(id=None, internal=None)) + + # Registering with the un-normalised identity must attach the + # channel to the already-added port, not create a second one. + assert manager.get_port_id(channel_id) == canonical + assert manager.get_port(canonical).get_channels() == {channel_id} + assert manager.get_all_data_channel_ids() == {channel_id} + + def test_add_input_port_twice_keeps_the_first_worker_port(self, manager): + port_id = PortIdentity(0, False) + first_schema = Schema(raw_schema={"x": "INTEGER"}) + manager.add_input_port(port_id, first_schema, [], []) + manager.register_input(_channel("upstream"), port_id) + + # A re-add (e.g. a second AddInputChannel for the same port) must + # not wipe the channels already registered against the port. + manager.add_input_port(port_id, Schema(raw_schema={"y": "STRING"}), [], []) + + port = manager.get_port(port_id) + assert port.get_schema() == first_schema + assert len(port.get_channels()) == 1 + + +class TestInputPortMaterializationReaderThreads: + @pytest.fixture + def manager(self): + return InputManager(worker_id=WORKER_ID, input_queue=MagicMock()) + + def test_reader_runnables_are_built_per_uri_and_exposed(self, manager): + port_id = PortIdentity(0, False) + partitioning_a = MagicMock(name="partitioning-a") + partitioning_b = MagicMock(name="partitioning-b") + + with patch.object( + input_manager_module, "InputPortMaterializationReaderRunnable" + ) as runnable_cls: + manager.add_input_port( + port_id, + Schema(raw_schema={"x": "INTEGER"}), + ["vfs:///a", "vfs:///b"], + [partitioning_a, partitioning_b], + ) + + # One runnable per (uri, partitioning) pair, wired to this + # worker's shared input queue. + assert runnable_cls.call_count == 2 + assert [call.kwargs["uri"] for call in runnable_cls.call_args_list] == [ + "vfs:///a", + "vfs:///b", + ] + assert [ + call.kwargs["partitioning"] for call in runnable_cls.call_args_list + ] == [ + partitioning_a, + partitioning_b, + ] + assert runnable_cls.call_args.kwargs["worker_actor_id"] == ActorVirtualIdentity( + WORKER_ID + ) + assert runnable_cls.call_args.kwargs["queue"] is manager._input_queue + + registered = manager.get_input_port_mat_reader_threads() + assert list(registered) == [port_id] + assert len(registered[port_id]) == 2 + + def test_no_uris_registers_an_empty_runnable_list(self, manager): + port_id = PortIdentity(0, False) + + manager.add_input_port(port_id, Schema(raw_schema={"x": "INTEGER"}), [], []) + + assert manager.get_input_port_mat_reader_threads() == {port_id: []} + + def test_mismatched_uris_and_partitionings_are_rejected(self, manager): + with pytest.raises(AssertionError): + manager.set_up_input_port_mat_reader_threads( + PortIdentity(0, False), ["vfs:///a"], [] + ) + + def test_start_threads_skips_already_finished_readers(self, manager): + # A reader that already drained its materialized input must not be + # replayed: restarting it would re-emit every tuple plus a second + # StartChannel/EndChannel pair on the same channel. + pending = _reader_stub(finished=False, channel_name="pending") + drained = _reader_stub(finished=True, channel_name="drained") + manager._input_port_mat_reader_runnables[PortIdentity(0, False)] = [ + pending, + drained, + ] + + with patch.object(input_manager_module, "threading") as threading_mock: + manager.start_input_port_mat_reader_threads() + + threading_mock.Thread.assert_called_once_with( + target=pending.run, + daemon=True, + name=f"port_mat_reader_runnable_thread_{pending.channel_id}", + ) + # Daemon threads are started, never joined, by this call. + threading_mock.Thread.return_value.start.assert_called_once_with() + + def test_start_threads_covers_every_port(self, manager): + first = _reader_stub(finished=False, channel_name="first") + second = _reader_stub(finished=False, channel_name="second") + manager._input_port_mat_reader_runnables[PortIdentity(0, False)] = [first] + manager._input_port_mat_reader_runnables[PortIdentity(1, False)] = [second] + + with patch.object(input_manager_module, "threading") as threading_mock: + manager.start_input_port_mat_reader_threads() + + started_targets = { + call.kwargs["target"] for call in threading_mock.Thread.call_args_list + } + assert started_targets == {first.run, second.run} + + +class TestProcessDataPayload: + @pytest.fixture + def manager(self): + manager = InputManager(worker_id=WORKER_ID, input_queue=MagicMock()) + manager.add_input_port( + PortIdentity(0, False), Schema(raw_schema={"x": "INTEGER"}), [], [] + ) + manager.register_input(_channel("upstream"), PortIdentity(0, False)) + return manager + + def test_unknown_payload_type_is_rejected(self, manager): + channel_id = _channel("upstream") + + # An unrecognised DataPayload must fail loudly rather than be + # silently dropped -- the tuple stream would go missing otherwise. + with pytest.raises(NotImplementedError): + list(manager.process_data_payload(channel_id, MagicMock(name="payload"))) + + def test_unknown_payload_type_still_records_the_current_channel(self, manager): + channel_id = _channel("upstream") + + with pytest.raises(NotImplementedError): + list(manager.process_data_payload(channel_id, object())) + + assert manager._current_channel_id == channel_id + + def test_state_frame_envelope_is_forwarded_untouched(self, manager): + channel_id = _channel("upstream") + state_frame = StateFrame(frame=State({"i": 1}), loop_counter=3) + + outputs = list(manager.process_data_payload(channel_id, state_frame)) + + # The whole envelope is yielded (not just .frame) so the runtime can + # read loop_counter off it. + assert outputs == [state_frame] + assert outputs[0].loop_counter == 3 + + def test_data_frame_is_expanded_into_schema_bearing_tuples(self, manager): + channel_id = _channel("upstream") + schema = Schema(raw_schema={"x": "INTEGER"}) + table = pyarrow.Table.from_pydict( + {"x": [1, 2]}, schema=schema.as_arrow_schema() + ) + + outputs = list(manager.process_data_payload(channel_id, DataFrame(frame=table))) + + assert [t["x"] for t in outputs] == [1, 2] + assert all(t._schema == schema for t in outputs) + + +class TestPortCompletion: + def test_ports_complete_independently(self): + manager = InputManager(worker_id=WORKER_ID, input_queue=MagicMock()) + port_a, port_b = PortIdentity(0, False), PortIdentity(1, False) + channel_a, channel_b = _channel("a"), _channel("b") + for port_id, channel_id in ((port_a, channel_a), (port_b, channel_b)): + manager.add_input_port(port_id, Schema(raw_schema={"x": "INTEGER"}), [], []) + manager.register_input(channel_id, port_id) + + assert manager.all_ports_completed() is False + + manager.complete_current_port(channel_a) + assert manager.all_ports_completed() is False + + manager.complete_current_port(channel_b) + assert manager.all_ports_completed() is True + assert set(manager.get_all_channel_ids()) == {channel_a, channel_b} diff --git a/amber/src/test/python/core/models/test_tuple.py b/amber/src/test/python/core/models/test_tuple.py index b9d4219913..aff43faa88 100644 --- a/amber/src/test/python/core/models/test_tuple.py +++ b/amber/src/test/python/core/models/test_tuple.py @@ -17,6 +17,7 @@ import datetime import pandas +import pickle import pyarrow import pytest import numpy as np @@ -657,3 +658,179 @@ class TestTuple: assert len(tuples) == 1 tuple_ = tuples[0] assert tuple_["large_binary_field"] is None + + def test_binary_field_round_trips_through_arrow(self): + # cast_to_schema pickles a non-bytes BINARY field behind a + # "pickle " sentinel; the arrow field accessor must recognise + # that sentinel and hand the original object back. The two halves + # are asserted together so the sentinel layout cannot drift. + original = Tuple({"scores": [85, 94, 100]}) + original.finalize(Schema(raw_schema={"scores": "BINARY"})) + pickled = original["scores"] + assert pickled[:10] == b"pickle " + + arrow_schema = pyarrow.schema( + [ + pyarrow.field("scores", pyarrow.binary()), + pyarrow.field("raw", pyarrow.binary()), + pyarrow.field("empty", pyarrow.binary()), + ] + ) + arrow_table = pyarrow.Table.from_pydict( + {"scores": [pickled], "raw": [b"plain bytes"], "empty": [None]}, + schema=arrow_schema, + ) + + tuples = [ + Tuple({name: field_accessor for name in arrow_table.column_names}) + for field_accessor in ArrowTableTupleProvider(arrow_table) + ] + + assert len(tuples) == 1 + assert tuples[0]["scores"] == [85, 94, 100] + # Bytes that were never pickled must survive untouched, and a null + # binary must not be fed to pickle.loads. + assert tuples[0]["raw"] == b"plain bytes" + assert tuples[0]["empty"] is None + + def test_binary_field_starting_with_pickle_bytes_is_unpickled(self): + # Sanity check on the sentinel itself: only the payload after the + # 10-byte header is unpickled. + payload = pickle.dumps({"k": "v"}) + arrow_table = pyarrow.Table.from_pydict( + {"blob": [b"pickle " + payload]}, + schema=pyarrow.schema([pyarrow.field("blob", pyarrow.binary())]), + ) + tuples = [ + Tuple({name: field_accessor for name in arrow_table.column_names}) + for field_accessor in ArrowTableTupleProvider(arrow_table) + ] + assert tuples[0]["blob"] == {"k": "v"} + + def test_get_serialized_field_converts_large_binary_to_uri(self): + from core.models.type.large_binary import largebinary + + schema = Schema( + raw_schema={"blob": "LARGE_BINARY", "name": "STRING", "size": "INTEGER"} + ) + blob = largebinary("s3://test-bucket/path/to/object") + tuple_ = Tuple({"blob": blob, "name": "obj", "size": 3}, schema=schema) + + # Only the LARGE_BINARY field is flattened to its URI; the arrow + # writer stores that string, not the handle object. + assert tuple_.get_serialized_field("blob") == "s3://test-bucket/path/to/object" + assert isinstance(tuple_.get_serialized_field("blob"), str) + assert tuple_.get_serialized_field("name") == "obj" + assert tuple_.get_serialized_field("size") == 3 + # The tuple itself keeps the handle. + assert tuple_["blob"] is blob + + def test_get_serialized_field_without_schema_returns_value_as_is(self): + from core.models.type.large_binary import largebinary + + blob = largebinary("s3://test-bucket/path/to/object") + tuple_ = Tuple({"blob": blob}) + + # With no schema there is no LARGE_BINARY declaration to act on. + assert tuple_.get_serialized_field("blob") is blob + + def test_cast_to_schema_skips_fields_absent_from_the_schema(self): + # A field the schema does not declare makes get_attr_type raise. + # That failure must be contained to the offending field: the cast + # loop logs it and keeps going, so later fields are still coerced + # (validate_schema is what ultimately rejects the stray field). + messages = [] + handler_id = logger.add(messages.append, level="WARNING") + try: + tuple_ = Tuple({"stray": 7.0, "weight": 119.0}) + tuple_.cast_to_schema(Schema(raw_schema={"weight": "INTEGER"})) + finally: + logger.remove(handler_id) + + assert tuple_["stray"] == 7.0 + assert type(tuple_["stray"]) is float + assert tuple_["weight"] == 119 + assert type(tuple_["weight"]) is int + assert any("stray" in str(message) for message in messages) + + def test_validate_schema_rejects_a_single_missing_field(self): + tuple_ = Tuple({"x": 1}) + + with pytest.raises(KeyError) as exc_info: + tuple_.validate_schema(Schema(raw_schema={"x": "INTEGER", "y": "STRING"})) + + message = str(exc_info.value) + assert "'y' is expected but missing" in message + assert "fields" not in message + + def test_validate_schema_rejects_several_missing_fields(self): + tuple_ = Tuple({"x": 1}) + + with pytest.raises(KeyError) as exc_info: + tuple_.validate_schema( + Schema(raw_schema={"x": "INTEGER", "y": "STRING", "z": "STRING"}) + ) + + message = str(exc_info.value) + assert "are expected but missing" in message + assert "'y'" in message and "'z'" in message + + def test_validate_schema_rejects_a_single_unexpected_field(self): + tuple_ = Tuple({"x": 1, "extra": "e"}) + + with pytest.raises(KeyError) as exc_info: + tuple_.validate_schema(Schema(raw_schema={"x": "INTEGER"})) + + message = str(exc_info.value) + assert "an unexpected field: 'extra'" in message + + def test_validate_schema_rejects_several_unexpected_fields(self): + tuple_ = Tuple({"x": 1, "extra": "e", "another": "a"}) + + with pytest.raises(KeyError) as exc_info: + tuple_.validate_schema(Schema(raw_schema={"x": "INTEGER"})) + + message = str(exc_info.value) + assert "unexpected fields" in message + assert "'extra'" in message and "'another'" in message + + def test_validate_schema_reports_missing_fields_before_unexpected_ones(self): + # Both problems at once: the missing-field error wins, because a + # renamed field is far more often a missing field than a stray one. + tuple_ = Tuple({"extra": "e"}) + + with pytest.raises(KeyError) as exc_info: + tuple_.validate_schema(Schema(raw_schema={"x": "INTEGER"})) + + assert "expected but missing" in str(exc_info.value) + + def test_tuple_iter_yields_field_values_in_order(self, target_tuple): + assert list(target_tuple) == [1, "a"] + assert tuple(target_tuple) == target_tuple.get_fields() + + def test_tuple_iter_resolves_lazy_accessors(self): + def field_accessor(field_name): + return chr(96 + int(field_name)) + + tuple_ = Tuple({"1": field_accessor, "3": field_accessor}) + + assert list(tuple_) == ["a", "c"] + + def test_tuple_contains_matches_field_names_not_values(self, target_tuple): + assert "x" in target_tuple + assert "y" in target_tuple + assert "z" not in target_tuple + # membership is by field name; a field *value* is not "in" the tuple + assert 1 not in target_tuple + assert "a" not in target_tuple + + def test_hash_treats_none_fields_as_zero(self): + # Mirrors java.util.Objects.hash: a null field contributes 0 to the + # rolling 31-based accumulator. Expected values are computed from + # that rule (1*31+0 = 31; 31*31+0 = 961). + schema = Schema(raw_schema={"a": "INTEGER", "b": "STRING"}) + + assert hash(Tuple({"a": None, "b": None}, schema)) == 961 + # Only the null field short-circuits; the other is hashed normally + # (java_hash_bytes("a") == 97, so 31*31 + 97 == 1058). + assert hash(Tuple({"a": None, "b": "a"}, schema)) == 1058 diff --git a/amber/src/test/python/core/storage/model/test_virtual_document.py b/amber/src/test/python/core/storage/model/test_virtual_document.py new file mode 100644 index 0000000000..bb90a445e3 --- /dev/null +++ b/amber/src/test/python/core/storage/model/test_virtual_document.py @@ -0,0 +1,104 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import pytest + +from core.storage.model.readonly_virtual_document import ReadonlyVirtualDocument +from core.storage.model.virtual_document import VirtualDocument + + +class _MinimalDocument(VirtualDocument[int]): + """The smallest concrete document a backend can supply. + + ``clear`` is the only abstract member of VirtualDocument, so an + implementation that supports nothing else must still be constructible + -- every other member falls back to the base "not implemented" + behaviour. Delegating to ``super().clear()`` also pins that the base + ``clear`` is a no-op rather than a raiser. + """ + + def __init__(self): + self.cleared = False + + def clear(self) -> None: + self.cleared = True + return super().clear() + + +class TestVirtualDocumentDefaults: + """VirtualDocument deliberately gives every read/write method a default + body that raises, so a backend that cannot reasonably support an + operation simply does not override it. These tests pin that each + default raises NotImplementedError (rather than silently returning + None) and names the method that is missing, which is what a caller + sees when it reaches for an unsupported operation. + """ + + @pytest.fixture + def document(self): + return _MinimalDocument() + + @pytest.mark.parametrize( + ("method_name", "args"), + [ + ("get_uri", ()), + ("get_item", (0,)), + ("get", ()), + ("get_range", (0, 1)), + ("get_after", (0,)), + ("get_count", ()), + ("writer", ("writer-1",)), + ], + ) + def test_unsupported_operations_raise_with_the_method_name( + self, document, method_name, args + ): + with pytest.raises(NotImplementedError) as exc_info: + getattr(document, method_name)(*args) + + assert str(exc_info.value) == f"{method_name} method is not implemented" + + def test_read_defaults_are_not_lazy_generators(self, document): + # get / get_range / get_after are declared as Iterator-returning. + # If their bodies were generators, calling them would hand back an + # un-started generator and the failure would surface far from the + # call site. They must raise eagerly instead. + for call in ( + lambda: document.get(), + lambda: document.get_range(0, 1), + lambda: document.get_after(0), + ): + with pytest.raises(NotImplementedError): + call() + + def test_clear_is_the_only_abstract_member(self, document): + # A backend only has to implement clear(); everything else is + # optional. The base clear() body itself is a no-op. + assert VirtualDocument.__abstractmethods__ == frozenset({"clear"}) + + document.clear() + + assert document.cleared is True + + def test_cannot_instantiate_without_clear(self): + with pytest.raises(TypeError, match="clear"): + VirtualDocument() + + def test_is_a_readonly_virtual_document(self, document): + # VirtualDocument extends the read-only abstraction, so anything + # accepting a ReadonlyVirtualDocument also accepts a writable one. + assert isinstance(document, ReadonlyVirtualDocument) diff --git a/amber/src/test/python/core/storage/runnables/test_input_port_materialization_reader_runnable.py b/amber/src/test/python/core/storage/runnables/test_input_port_materialization_reader_runnable.py index 288ee6374b..56069b3163 100644 --- a/amber/src/test/python/core/storage/runnables/test_input_port_materialization_reader_runnable.py +++ b/amber/src/test/python/core/storage/runnables/test_input_port_materialization_reader_runnable.py @@ -18,9 +18,10 @@ from unittest.mock import MagicMock, patch import pytest +from loguru import logger -from core.models import State, StateFrame -from core.models.internal_queue import DataElement +from core.models import DataFrame, State, StateFrame, Tuple +from core.models.internal_queue import DataElement, ECMElement from core.models.schema import Schema from core.storage.runnables.input_port_materialization_reader_runnable import ( InputPortMaterializationReaderRunnable, @@ -30,6 +31,52 @@ from proto.org.apache.texera.amber.core import ( ChannelIdentity, ) +DOCUMENT_FACTORY = ( + "core.storage.runnables.input_port_materialization_reader_runnable.DocumentFactory" +) + + +def _build_reader(worker_actor_id, schema=None): + """Build a reader without running __init__ (which would resolve a real + partitioner and a real storage URI). Mirrors the fixture in + TestRunStateReadingBlock, but parameterised so the tuple-loop tests can + drive the partitioner themselves. + """ + reader = InputPortMaterializationReaderRunnable.__new__( + InputPortMaterializationReaderRunnable + ) + reader.uri = "vfs:///wf/0/exec/0/result/op-a" + reader.worker_actor_id = worker_actor_id + reader.tuple_schema = ( + schema if schema is not None else Schema(raw_schema={"x": "INTEGER"}) + ) + reader._stopped = False + reader._finished = False + reader.channel_id = ChannelIdentity( + worker_actor_id, worker_actor_id, is_control=False + ) + reader.queue = MagicMock() + reader.partitioner = MagicMock() + reader.partitioner.flush.return_value = [] + return reader + + +def _emitted_data_frames(reader): + return [ + call.args[0].payload + for call in reader.queue.put.call_args_list + if isinstance(call.args[0], DataElement) + and isinstance(call.args[0].payload, DataFrame) + ] + + +def _emitted_ecm_names(reader): + return [ + call.args[0].payload.id.id + for call in reader.queue.put.call_args_list + if isinstance(call.args[0], ECMElement) + ] + class TestRunStateReadingBlock: """Cover the state-reading block in run() that opens the state @@ -108,3 +155,195 @@ class TestRunStateReadingBlock: assert [sf.payload.loop_counter for sf in state_frames] == [0, 1] assert [sf.payload.loop_start_id for sf in state_frames] == ["loop-a", "loop-b"] assert runnable._finished is True + + +class TestTupleToBatchWithFilter: + """The reader reuses the *output*-side partitioner to decide which + tuples this worker would have received, then keeps only its own + share. Batches addressed to the other (hypothetical) downstream + workers must be dropped, otherwise a shuffled materialized port would + replay every partition into every worker. + """ + + @pytest.fixture + def me(self): + return ActorVirtualIdentity(name="me") + + def test_only_batches_addressed_to_this_worker_are_yielded(self, me): + other = ActorVirtualIdentity(name="other") + reader = _build_reader(me) + reader.partitioner.add_tuple_to_batch.return_value = [ + (other, [Tuple({"x": 9})]), + (me, [Tuple({"x": 1}), Tuple({"x": 2})]), + (other, [Tuple({"x": 8})]), + ] + + frames = list(reader.tuple_to_batch_with_filter(Tuple({"x": 0}))) + + assert len(frames) == 1 + assert frames[0].frame.to_pydict() == {"x": [1, 2]} + + def test_no_batch_for_this_worker_yields_nothing(self, me): + other = ActorVirtualIdentity(name="other") + reader = _build_reader(me) + reader.partitioner.add_tuple_to_batch.return_value = [ + (other, [Tuple({"x": 9})]), + ] + + assert list(reader.tuple_to_batch_with_filter(Tuple({"x": 0}))) == [] + + def test_tuples_to_data_frame_projects_the_schema_column_order(self, me): + # Column order and arrow types come from the tuple schema, not from + # the tuple's own field order -- the downstream arrow reader relies + # on that alignment. + schema = Schema(raw_schema={"x": "INTEGER", "y": "STRING"}) + reader = _build_reader(me, schema=schema) + + frame = reader.tuples_to_data_frame( + [Tuple({"y": "a", "x": 1}), Tuple({"y": "b", "x": 2})] + ) + + assert frame.frame.column_names == ["x", "y"] + assert frame.frame.to_pydict() == {"x": [1, 2], "y": ["a", "b"]} + assert frame.frame.schema == schema.as_arrow_schema() + + +class TestRunTupleLoop: + """Cover the tuple-replay half of run(): each stored tuple is cast to + the port schema, routed through the partitioner filter, and emitted as + a DataFrame between the StartChannel and EndChannel ECMs. + """ + + @pytest.fixture + def me(self): + return ActorVirtualIdentity(name="me") + + @staticmethod + def _run_with_documents(reader, tuples): + result_doc = MagicMock() + result_doc.get.return_value = iter(tuples) + state_doc = MagicMock() + state_doc.get.return_value = iter([]) # no materialized states + with patch(DOCUMENT_FACTORY) as mock_factory: + mock_factory.open_document.side_effect = [ + (result_doc, reader.tuple_schema), + (state_doc, None), + ] + reader.run() + return result_doc + + def test_tuples_are_cast_to_the_port_schema_before_emission(self, me): + reader = _build_reader(me) + # A materialized INT column can come back as a float (pandas + # promotes null-holding int columns to float64). Without the + # cast_to_schema() call the arrow conversion below would not + # produce ints. + reader.partitioner.add_tuple_to_batch.side_effect = lambda tup: [(me, [tup])] + + self._run_with_documents(reader, [Tuple({"x": 1.0}), Tuple({"x": 2})]) + + frames = _emitted_data_frames(reader) + assert [frame.frame.to_pydict() for frame in frames] == [ + {"x": [1]}, + {"x": [2]}, + ] + assert all( + isinstance(value, int) + for frame in frames + for value in frame.frame.column("x").to_pylist() + ) + assert reader.finished() is True + + def test_channel_is_bracketed_by_start_and_end_ecms(self, me): + reader = _build_reader(me) + # Make the partitioner behave like the real one: flush hands back + # the ECM so it reaches the queue. + reader.partitioner.flush.side_effect = lambda receiver, ecm: [ecm] + reader.partitioner.add_tuple_to_batch.side_effect = lambda tup: [(me, [tup])] + + self._run_with_documents(reader, [Tuple({"x": 1})]) + + assert _emitted_ecm_names(reader) == ["StartChannel", "EndChannel"] + # The data must land between the two ECMs. + payload_kinds = [ + type(call.args[0]).__name__ for call in reader.queue.put.call_args_list + ] + assert payload_kinds == ["ECMElement", "DataElement", "ECMElement"] + + def test_stop_halts_replay_but_still_closes_the_channel(self, me): + reader = _build_reader(me) + reader.partitioner.flush.side_effect = lambda receiver, ecm: [ecm] + reader.partitioner.add_tuple_to_batch.side_effect = lambda tup: [(me, [tup])] + + reader.stop() + self._run_with_documents(reader, [Tuple({"x": 1}), Tuple({"x": 2})]) + + # No tuple is replayed once stopped ... + assert _emitted_data_frames(reader) == [] + reader.partitioner.add_tuple_to_batch.assert_not_called() + # ... but the channel is still closed, so the downstream port + # alignment does not hang waiting for an EndChannel. + assert _emitted_ecm_names(reader) == ["StartChannel", "EndChannel"] + assert reader.finished() is True + + def test_stop_sets_the_flag_before_run(self, me): + reader = _build_reader(me) + assert reader._stopped is False + + reader.stop() + + assert reader._stopped is True + + +class TestRunErrorHandling: + """run() executes on a detached daemon thread, so it must swallow and + log failures instead of propagating them into an unhandled thread + exception. A reader that failed must not report itself as finished -- + otherwise start_input_port_mat_reader_threads would skip retrying it. + """ + + @pytest.fixture + def me(self): + return ActorVirtualIdentity(name="me") + + def test_open_document_failure_is_logged_and_leaves_reader_unfinished(self, me): + reader = _build_reader(me) + messages = [] + handler_id = logger.add(messages.append, level="ERROR") + try: + with patch(DOCUMENT_FACTORY) as mock_factory: + mock_factory.open_document.side_effect = RuntimeError( + "catalog unreachable" + ) + reader.run() # must not raise + finally: + logger.remove(handler_id) + + assert reader.finished() is False + reader.queue.put.assert_not_called() + assert any("catalog unreachable" in str(message) for message in messages) + + def test_failure_midway_leaves_the_reader_unfinished(self, me): + reader = _build_reader(me) + reader.partitioner.add_tuple_to_batch.side_effect = RuntimeError( + "bad partition" + ) + result_doc = MagicMock() + result_doc.get.return_value = iter([Tuple({"x": 1})]) + state_doc = MagicMock() + state_doc.get.return_value = iter([]) + + handler_id = logger.add(lambda _: None, level="ERROR") + try: + with patch(DOCUMENT_FACTORY) as mock_factory: + mock_factory.open_document.side_effect = [ + (result_doc, reader.tuple_schema), + (state_doc, None), + ] + reader.run() + finally: + logger.remove(handler_id) + + # The EndChannel is never reached, so the reader stays unfinished. + assert reader.finished() is False + assert _emitted_ecm_names(reader) == [] diff --git a/amber/src/test/python/pytexera/udf/test_udf_operator.py b/amber/src/test/python/pytexera/udf/test_udf_operator.py index b96199d450..a04e52b3a1 100644 --- a/amber/src/test/python/pytexera/udf/test_udf_operator.py +++ b/amber/src/test/python/pytexera/udf/test_udf_operator.py @@ -18,10 +18,12 @@ import datetime from typing import Iterator, Optional +import pandas import pytest import pytexera.udf.udf_operator as udf_operator from pytexera import AttributeType, Tuple, TupleLike, UDFOperatorV2 +from pytexera import UDFBatchOperator, UDFSourceOperator, UDFTableOperator from pytexera.udf.udf_operator import _UiParameterSupport @@ -143,6 +145,45 @@ class SuperOpenConflictingParameterOperator(UDFOperatorV2): yield tuple_ +class InvalidValueParameterOperator(UDFOperatorV2): + def _texera_injected_ui_parameters(self): + return {"count": "not-a-number"} + + def open(self): + self.count_parameter = self.UiParameter("count", AttributeType.INT) + + def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: + yield tuple_ + + +class NoOpenParameterSupport(_UiParameterSupport): + """A _UiParameterSupport subclass that defines no open() hook. + + __init_subclass__ has nothing to wrap here and must bail out instead + of installing a wrapper around a non-existent attribute. + """ + + +class BaseHookTupleOperator(UDFOperatorV2): + def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: + yield from super().process_tuple(tuple_, port) + + +class BaseHookSourceOperator(UDFSourceOperator): + def produce(self): + yield from super().produce() + + +class BaseHookTableOperator(UDFTableOperator): + def process_table(self, table, port: int): + yield from super().process_table(table, port) + + +class BaseHookBatchOperator(UDFBatchOperator): + def process_batch(self, batch, port: int): + yield from super().process_batch(batch, port) + + class TestUiParameterSupport: def test_injected_values_are_applied_before_open(self): operator = InjectedParametersOperator() @@ -351,3 +392,111 @@ class TestUiParameterSupport: first_operator._ui_parameter_injected_values is not second_operator._ui_parameter_injected_values ) + + def test_subclass_without_open_is_left_unwrapped(self): + # Not every _UiParameterSupport subclass declares open(); the + # wrapper installation must be skipped rather than fabricating an + # open() attribute on such a class. + assert getattr(NoOpenParameterSupport, "open", None) is None + + support = NoOpenParameterSupport() + parameter = support.UiParameter("absent", AttributeType.STRING) + + # UiParameter still works: the state is lazily initialised, and an + # un-injected name resolves to None. + assert parameter.value is None + assert support._ui_parameter_injected_values == {} + assert support._ui_parameter_name_types == {"absent": AttributeType.STRING} + + @pytest.mark.parametrize( + ("raw_value", "attr_type"), + [ + ("not-a-number", AttributeType.INT), + ("not-a-number", AttributeType.LONG), + ("12.5.6", AttributeType.DOUBLE), + ("not-a-timestamp", AttributeType.TIMESTAMP), + ], + ) + def test_parse_unparseable_value_raises_value_error(self, raw_value, attr_type): + with pytest.raises(ValueError) as exc_info: + _UiParameterSupport._parse(raw_value, attr_type) + + message = str(exc_info.value) + assert f"Failed to parse UiParameter value {raw_value!r}" in message + assert f"as {attr_type.name}" in message + assert f"valid {attr_type.name.lower()} value" in message + # The underlying parser error is chained, not swallowed. + assert exc_info.value.__cause__ is not None + + def test_unparseable_injected_value_fails_at_open(self): + operator = InvalidValueParameterOperator() + + with pytest.raises(ValueError, match="Failed to parse UiParameter value"): + operator.open() + + def test_wrapped_open_resets_the_in_progress_flag_after_a_failure(self): + # The flag is cleared in a finally block, so a second open() after + # a failing one must re-apply the injected values (rather than take + # the "already in progress" shortcut and skip injection). + operator = InvalidValueParameterOperator() + with pytest.raises(ValueError): + operator.open() + + assert operator._ui_parameter_open_in_progress is False + + with pytest.raises(ValueError): + operator.open() + + +class TestBaseOperatorHooks: + """Every UDF base class ships a default generator body so a user + operator can call super() (or leave an optional hook alone) and simply + produce nothing. Each default must yield exactly one None, which the + runtime filters out -- returning a non-generator would break the + caller's `for ... in` loop instead. + """ + + def test_tuple_operator_defaults_yield_nothing_useful(self): + operator = BaseHookTupleOperator() + + assert operator.open() is None + assert list(operator.process_tuple(Tuple({"x": 1}), 0)) == [None] + assert list(operator.on_finish(0)) == [None] + assert operator.close() is None + + def test_source_operator_defaults_yield_nothing_useful(self): + operator = BaseHookSourceOperator() + + assert operator.open() is None + assert list(operator.produce()) == [None] + assert operator.close() is None + + def test_table_operator_defaults_yield_nothing_useful(self): + operator = BaseHookTableOperator() + table = pandas.DataFrame({"x": [1, 2]}) + + assert operator.open() is None + assert list(operator.process_table(table, 0)) == [None] + assert operator.close() is None + + def test_batch_operator_defaults_yield_nothing_useful(self): + operator = BaseHookBatchOperator() + batch = pandas.DataFrame({"x": [1, 2]}) + + assert operator.open() is None + assert list(operator.process_batch(batch, 0)) == [None] + assert operator.close() is None + + @pytest.mark.parametrize( + ("operator_class", "hook_name"), + [ + (UDFOperatorV2, "process_tuple"), + (UDFSourceOperator, "produce"), + (UDFTableOperator, "process_table"), + (UDFBatchOperator, "process_batch"), + ], + ) + def test_processing_hook_is_abstract(self, operator_class, hook_name): + # The default body exists for super() delegation only -- a user + # operator must still implement the hook. + assert hook_name in operator_class.__abstractmethods__
