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-8769-fc62c09a27bf5e9aa73027ce9b37fec5fb5bd018 in repository https://gitbox.apache.org/repos/asf/texera.git
commit 1d8351b1bbe1226ccc74e832c3017626f734920d Author: Suryaa Charan <[email protected]> AuthorDate: Thu Oct 15 18:22:52 2026 +0000 feat(pyamber): clear errors for return-instead-of-yield and output mismatches (#8769) ### What changes were proposed in this PR? Clear messages for the common Python UDF mistakes that show up while the workflow runs: - **`return` instead of `yield`** in `process_tuple`, `process_batch`, `process_table` or `produce` now says `process_table must yield results, not return them. Replace return x with yield x.` Before, users saw `pampy.MatchError: '_' not provided` or `'NoneType' object is not iterable`. Returning a list still works. - **`__init__` without `super().__init__()`** on a table or batch UDF no longer crashes with `'_TableOperator__table_data'`. The buffers are created on first use, the same way `Operator` already creates its template decoder. - **A DataFrame in a non-binary column** now says which column, and to yield the DataFrame itself or make the column binary. A DataFrame in a BINARY column is pickled as before. - **Output columns that don't match** keep the original message and point to "Retain input columns" and "Extra output column(s)" (or "Columns" for a 1-out Python UDF). Before and after, in the Console: | Before | After | |---|---| |<img width="800" height="288" alt="runtime-before" src="https://github.com/user-attachments/assets/7f854ffe-f5eb-4013-8fe4-9eb3a9e740ff"; /> | <img width="800" height="288" alt="runtime-after" src="https://github.com/user-attachments/assets/fbaec933-683f-4dd2-b603-5d988b4f338d"; /> | ### Any related issues, documentation, discussions? Part of #8736. Split out of #8737. ### How was this PR tested? New tests in `test_operator.py` (each `return` case, a returned list still working, table and batch UDFs with and without `super().__init__()`) and `test_data_processor.py` (`process_tuple` and `on_finish` returning a value, a DataFrame in BINARY, LARGE_BINARY and STRING columns, a missing or misspelled column next to a DataFrame, a wrong column type). ``` cd amber pytest -m "not integration" src/test/python/core/models/test_operator.py src/test/python/core/runnables/test_data_processor.py pytest -m "not integration" ruff check src/main/python src/test/python && ruff format --check src/main/python src/test/python ``` The screenshots are from a local Texera built from this branch. ### Was this PR authored or co-authored using generative AI tooling? Generated-by: Claude Code (Claude Opus 5.5) --------- Co-authored-by: suryaacharan <[email protected]> Co-authored-by: Claude Opus 5.5 <[email protected]> --- amber/src/main/python/core/models/operator.py | 60 +++++-- .../main/python/core/runnables/data_processor.py | 55 ++++++- amber/src/test/python/core/models/test_operator.py | 111 +++++++++++++ .../python/core/runnables/test_data_processor.py | 175 +++++++++++++++++++++ 4 files changed, 380 insertions(+), 21 deletions(-) diff --git a/amber/src/main/python/core/models/operator.py b/amber/src/main/python/core/models/operator.py index 4e65fbb2ab..bbba116da5 100644 --- a/amber/src/main/python/core/models/operator.py +++ b/amber/src/main/python/core/models/operator.py @@ -29,6 +29,22 @@ from .table import all_output_to_tuple import base64 +def require_yielded(result, method_name: str): + """ + Catch the common case of a user method such as process_table using + `return` instead of `yield`: it rejects a returned None, Tuple, dict or + DataFrame, which would otherwise fail later with a confusing error. + Other values (e.g. a list, or a pandas Series) are passed through + unchecked. + """ + if result is None or isinstance(result, (Tuple, dict, pandas.DataFrame)): + raise TypeError( + f"{method_name} must `yield` results, not `return` them. " + "Replace `return x` with `yield x`." + ) + return result + + class Operator(ABC): """ Abstract base class for all operators. @@ -179,7 +195,7 @@ class SourceOperator(TupleOperatorV2): @overrides.final def on_finish(self, port: int) -> Iterator[Optional[TupleLike]]: # TODO: change on_finish to output Iterator[Union[TupleLike, TableLike, None]] - for i in self.produce(): + for i in require_yielded(self.produce(), "produce"): yield from all_output_to_tuple(i) @overrides.final @@ -197,7 +213,6 @@ class BatchOperator(TupleOperatorV2): def __init__(self): super().__init__() - self.__batch_data: MutableMapping[int, List[Tuple]] = defaultdict(list) self._validate_batch_size(self.BATCH_SIZE) @staticmethod @@ -210,26 +225,35 @@ class BatchOperator(TupleOperatorV2): if value <= 0: raise ValueError("BATCH_SIZE should be positive.") + def _get_batch_data(self) -> MutableMapping[int, List[Tuple]]: + # Created on first use (not in __init__) so a subclass whose __init__ + # skips super().__init__() still works, like _get_template_decoder. + if not hasattr(self, "_BatchOperator__batch_data"): + self._validate_batch_size(self.BATCH_SIZE) + self.__batch_data = defaultdict(list) + return self.__batch_data + @overrides.final def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: - self.__batch_data[port].append(tuple_) - if ( - self.BATCH_SIZE is not None - and len(self.__batch_data[port]) >= self.BATCH_SIZE - ): + batch_data = self._get_batch_data() + batch_data[port].append(tuple_) + if self.BATCH_SIZE is not None and len(batch_data[port]) >= self.BATCH_SIZE: yield from self._process_batch(port) @overrides.final def _process_batch(self, port: int) -> Iterator[Optional[BatchLike]]: + batch_data = self._get_batch_data() batch = Batch( pandas.DataFrame( [ - self.__batch_data[port].pop(0).as_series() - for _ in range(min(len(self.__batch_data[port]), self.BATCH_SIZE)) + batch_data[port].pop(0).as_series() + for _ in range(min(len(batch_data[port]), self.BATCH_SIZE)) ] ) ) - for output_batch in self.process_batch(batch, port): + for output_batch in require_yielded( + self.process_batch(batch, port), "process_batch" + ): if output_batch is not None: if isinstance(output_batch, pandas.DataFrame): # TODO: integrate into Batch as a helper function. @@ -242,7 +266,7 @@ class BatchOperator(TupleOperatorV2): @overrides.final def on_finish(self, port: int) -> Iterator[Optional[BatchLike]]: - while len(self.__batch_data[port]) != 0: + while len(self._get_batch_data()[port]) != 0: yield from self._process_batch(port) @abstractmethod @@ -268,16 +292,22 @@ class TableOperator(TupleOperatorV2): def __init__(self): super().__init__() self._Operator__internal_is_source: bool = False - self.__table_data: Mapping[int, List[Tuple]] = defaultdict(list) + + def _get_table_data(self) -> Mapping[int, List[Tuple]]: + # Created on first use (not in __init__) so a subclass whose __init__ + # skips super().__init__() still works, like _get_template_decoder. + if not hasattr(self, "_TableOperator__table_data"): + self.__table_data = defaultdict(list) + return self.__table_data @overrides.final def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: - self.__table_data[port].append(tuple_) + self._get_table_data()[port].append(tuple_) yield def on_finish(self, port: int) -> Iterator[Optional[TableLike]]: - table = Table(self.__table_data[port]) - yield from self.process_table(table, port) + table = Table(self._get_table_data()[port]) + yield from require_yielded(self.process_table(table, port), "process_table") @abstractmethod def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]: diff --git a/amber/src/main/python/core/runnables/data_processor.py b/amber/src/main/python/core/runnables/data_processor.py index eaee8660cd..4466cc6adb 100644 --- a/amber/src/main/python/core/runnables/data_processor.py +++ b/amber/src/main/python/core/runnables/data_processor.py @@ -16,12 +16,15 @@ # under the License. from contextlib import contextmanager + +import pandas from threading import Event from typing import Iterator, Optional from core.architecture.managers import Context -from core.models import State, TupleLike +from core.models import AttributeType, Schema, State, TupleLike from core.models.internal_marker import EndChannel, PortMarker, StartChannel +from core.models.operator import require_yielded from core.models.table import all_output_to_tuple from core.util import Stoppable from core.util.console_message.replace_print import replace_print @@ -76,7 +79,9 @@ class DataProcessor(Runnable, Stoppable): # Flush the state to MainLoop before producing tuples so the # state and the tuple stream don't share a single switch. self._switch_context() - self._set_output_tuple(executor.on_finish(port_id)) + self._set_output_tuple( + require_yielded(executor.on_finish(port_id), "on_finish") + ) def process_state(self, state: State) -> None: """ @@ -94,7 +99,11 @@ class DataProcessor(Runnable, Stoppable): while not finished_current.is_set(): with self._executor_session() as (executor, port_id): tuple_ = self._context.tuple_processing_manager.get_input_tuple() - self._set_output_tuple(executor.process_tuple(tuple_, port_id)) + self._set_output_tuple( + require_yielded( + executor.process_tuple(tuple_, port_id), "process_tuple" + ) + ) @contextmanager def _executor_session(self, marker_port_id: int | None = None): @@ -132,9 +141,14 @@ class DataProcessor(Runnable, Stoppable): # output could be a None, a TupleLike, or a TableLike. for output_tuple in all_output_to_tuple(output): if output_tuple is not None: - output_tuple.finalize( - self._context.output_manager.get_port().get_schema() - ) + schema = self._context.output_manager.get_port().get_schema() + try: + output_tuple.finalize(schema) + except (KeyError, TypeError) as error: + # Only explain on failure, so valid rows are not inspected. + raise self._explain_schema_mismatch( + output, schema, error + ) from error self._switch_context() self._context.tuple_processing_manager.current_output_tuple = ( output_tuple @@ -142,6 +156,35 @@ class DataProcessor(Runnable, Stoppable): self._switch_context() self._context.tuple_processing_manager.finished_current.set() + @staticmethod + def _explain_schema_mismatch(output, schema: Schema, error: Exception) -> Exception: + """ + Rebuild a schema check failure with a hint on how to fix the output, + keeping the original message. The schema check reports missing or + extra columns (KeyError) before wrong types (TypeError), so only a + TypeError can be caused by a DataFrame sitting in a column; a + DataFrame is valid in a BINARY column, where it is pickled. + """ + message = error.args[0] if error.args else str(error) + if isinstance(error, TypeError) and isinstance(output, dict): + columns = set(schema.get_attr_names()) + for key, value in output.items(): + if ( + isinstance(value, pandas.DataFrame) + and key in columns + and schema.get_attr_type(key) != AttributeType.BINARY + ): + return TypeError( + f"{message} A DataFrame was put in column '{key}'. To " + "output its rows, yield the DataFrame itself instead of " + "a dict; to keep it as one value, make the column binary." + ) + return type(error)( + f"{message} The yielded columns and their types must match the " + "operator's output columns: 'Retain input columns' and 'Extra output " + "column(s)' for a Python UDF, or 'Columns' for a 1-out Python UDF." + ) + def _set_output_state(self, output_state: State) -> None: """ Set the output state after processing by the executor. diff --git a/amber/src/test/python/core/models/test_operator.py b/amber/src/test/python/core/models/test_operator.py index bb18853350..6ca9ed5658 100644 --- a/amber/src/test/python/core/models/test_operator.py +++ b/amber/src/test/python/core/models/test_operator.py @@ -514,3 +514,114 @@ class TestBatchOperatorBatchAssembly: assert op.seen == [[[1], [2]]] # the port-1 tuple is not in here assert op.ports == [0] assert out == [Tuple({"batched": 2})] + + +class _TableWithoutSuperInit(TableOperator): + def __init__(self): + self.threshold = 1 # forgets super().__init__() + + def process_table(self, table, port): + yield table + + +class _BatchWithoutSuperInit(BatchOperator): + BATCH_SIZE = 2 + + def __init__(self): + self.seen = 0 # forgets super().__init__() + + def process_batch(self, batch, port): + yield batch + + +class _TableReturningDataFrame(TableOperator): + def process_table(self, table, port): + return table + + +class _TableReturningNone(TableOperator): + def process_table(self, table, port): + return None + + +class _TableReturningList(TableOperator): + def process_table(self, table, port): + return [{"x": 1}] + + +class _BatchReturningBatch(BatchOperator): + BATCH_SIZE = 1 + + def process_batch(self, batch, port): + return batch + + +class _SourceReturningDict(SourceOperator): + def produce(self): + return {"x": 1} + + +class TestCommonUDFMistakes: + """Mistakes UDF authors make should either work or say how to fix them.""" + + def test_table_operator_works_without_super_init(self): + op = _TableWithoutSuperInit() + list(op.process_tuple(Tuple({"x": 1}), 0)) + (table,) = list(op.on_finish(0)) + assert list(table["x"]) == [1] + + def test_batch_operator_works_without_super_init(self): + op = _BatchWithoutSuperInit() + list(op.process_tuple(Tuple({"x": 1}), 0)) + outputs = list(op.process_tuple(Tuple({"x": 2}), 0)) + assert [t["x"] for t in outputs] == [1, 2] + + @pytest.mark.parametrize( + "operator", [_TableReturningDataFrame, _TableReturningNone] + ) + def test_process_table_return_says_to_yield(self, operator): + op = operator() + list(op.process_tuple(Tuple({"x": 1}), 0)) + with pytest.raises(TypeError) as exc_info: + list(op.on_finish(0)) + assert "process_table must `yield` results, not `return` them" in str( + exc_info.value + ) + + def test_process_table_returning_a_list_still_works(self): + op = _TableReturningList() + assert list(op.on_finish(0)) == [{"x": 1}] + + def test_process_batch_return_says_to_yield(self): + op = _BatchReturningBatch() + with pytest.raises(TypeError) as exc_info: + list(op.process_tuple(Tuple({"x": 1}), 0)) + assert "process_batch must `yield` results" in str(exc_info.value) + + def test_produce_return_says_to_yield(self): + op = _SourceReturningDict() + with pytest.raises(TypeError) as exc_info: + list(op.on_finish(0)) + assert "produce must `yield` results" in str(exc_info.value) + + +class TestBufferCreatedOnlyByGetter: + def test_table_operator_has_no_buffer_until_first_tuple(self): + op = _TableWithoutSuperInit() + assert "_TableOperator__table_data" not in op.__dict__ + + class _WithSuper(TableOperator): + def process_table(self, table, port): + yield table + + assert "_TableOperator__table_data" not in _WithSuper().__dict__ + + def test_batch_operator_still_rejects_a_bad_batch_size_at_construction(self): + class _BadSize(BatchOperator): + BATCH_SIZE = 0 + + def process_batch(self, batch, port): + yield batch + + with pytest.raises(ValueError): + _BadSize() diff --git a/amber/src/test/python/core/runnables/test_data_processor.py b/amber/src/test/python/core/runnables/test_data_processor.py index e185872f2e..6b321dd299 100644 --- a/amber/src/test/python/core/runnables/test_data_processor.py +++ b/amber/src/test/python/core/runnables/test_data_processor.py @@ -290,3 +290,178 @@ class TestRunInvariant: assert "expected exactly one queued input" in str(excinfo.value) assert "marker=True" in str(excinfo.value) assert "tuple=True" in str(excinfo.value) + + +class _ReturningExecutor: + """process_tuple uses `return` instead of `yield`.""" + + def process_tuple(self, tuple_, port_id): + return tuple_ + + +class TestUDFOutputMistakes: + @pytest.fixture + def output_schema(self, context, monkeypatch): + schema = Schema(raw_schema={"x": "INTEGER"}) + monkeypatch.setattr( + context.output_manager, + "get_port", + lambda *args, **kwargs: SimpleNamespace(get_schema=lambda: schema), + ) + return schema + + @pytest.mark.timeout(2) + def test_process_tuple_return_says_to_yield(self, context, data_processor): + context.executor_manager.executor = _ReturningExecutor() + context.tuple_processing_manager.current_input_tuple = Tuple({"x": 1}) + + def switch_and_stop(): + data_processor.switch_calls += 1 + context.tuple_processing_manager.finished_current.set() + + data_processor._switch_context = switch_and_stop + data_processor.process_tuple() + + exc_info = context.exception_manager.get_exc_info() + assert exc_info[0] is TypeError + assert "process_tuple must `yield` results, not `return` them" in str( + exc_info[1] + ) + + def test_dict_of_dataframes_says_to_yield_the_dataframe( + self, data_processor, output_schema + ): + import pandas + + with pytest.raises(TypeError) as exc_info: + data_processor._set_output_tuple( + iter([{"x": pandas.DataFrame({"x": [1]})}]) + ) + assert "yield the DataFrame itself instead of a dict" in str(exc_info.value) + + def test_schema_mismatch_points_to_output_columns( + self, data_processor, output_schema + ): + with pytest.raises(KeyError) as exc_info: + data_processor._set_output_tuple(iter([{"y": 1}])) + message = str(exc_info.value) + assert "expected but missing" in message + assert "Extra output column(s)" in message + + +class _FinishReturningExecutor: + """on_finish uses `return` instead of `yield`.""" + + def produce_state_on_finish(self, port_id): + return None + + def on_finish(self, port_id): + return {"x": 1} + + +class TestUDFOutputMistakesFollowUp: + def _use_schema(self, context, monkeypatch, raw_schema): + schema = Schema(raw_schema=raw_schema) + monkeypatch.setattr( + context.output_manager, + "get_port", + lambda *args, **kwargs: SimpleNamespace(get_schema=lambda: schema), + ) + + def test_dict_holding_a_dataframe_into_a_binary_column_still_works( + self, context, data_processor, monkeypatch + ): + import pandas + + # BINARY columns pickle any value, so this is valid output on main. + self._use_schema(context, monkeypatch, {"result": "BINARY"}) + data_processor._set_output_tuple( + iter([{"result": pandas.DataFrame({"x": [1]})}]) + ) + assert context.tuple_processing_manager.finished_current.is_set() + + def test_wrong_column_type_points_to_output_columns( + self, context, data_processor, monkeypatch + ): + self._use_schema(context, monkeypatch, {"x": "INTEGER"}) + with pytest.raises(TypeError) as exc_info: + data_processor._set_output_tuple(iter([{"x": "not a number"}])) + message = str(exc_info.value) + assert "Unmatched type for field 'x'" in message + assert "Extra output column(s)" in message + + @pytest.mark.timeout(2) + def test_on_finish_return_says_to_yield(self, context, data_processor): + context.executor_manager.executor = _FinishReturningExecutor() + data_processor.process_internal_marker(EndChannel(0)) + exc_info = context.exception_manager.get_exc_info() + assert exc_info[0] is TypeError + assert "on_finish must `yield` results, not `return` them" in str(exc_info[1]) + + def test_misspelled_column_next_to_a_dataframe_gets_the_column_hint( + self, context, data_processor, monkeypatch + ): + import pandas + + # The DataFrame is fine (BINARY column); the column name is wrong. + self._use_schema(context, monkeypatch, {"result": "BINARY"}) + with pytest.raises(KeyError) as exc_info: + data_processor._set_output_tuple( + iter([{"res": pandas.DataFrame({"x": [1]})}]) + ) + message = str(exc_info.value) + assert "Extra output column(s)" in message + assert "dict of DataFrames" not in message + + def test_dataframe_in_a_non_binary_column_says_to_yield_the_dataframe( + self, context, data_processor, monkeypatch + ): + import pandas + + self._use_schema(context, monkeypatch, {"result": "STRING"}) + with pytest.raises(TypeError) as exc_info: + data_processor._set_output_tuple( + iter([{"result": pandas.DataFrame({"x": [1]})}]) + ) + assert "yield the DataFrame itself" in str(exc_info.value) + + def test_missing_column_is_reported_even_with_a_dataframe_in_another_column( + self, context, data_processor, monkeypatch + ): + import pandas + + # The real problem is the missing "count" column, not the DataFrame. + self._use_schema( + context, + monkeypatch, + {"name": "STRING", "count": "INTEGER", "stats": "STRING"}, + ) + with pytest.raises(KeyError) as exc_info: + data_processor._set_output_tuple( + iter([{"name": "a", "stats": pandas.DataFrame({"v": [1]})}]) + ) + message = str(exc_info.value) + assert "'count'" in message and "expected but missing" in message + assert "DataFrame" not in message.split("]")[-1] + + def test_dataframe_hint_names_the_column_and_keeps_the_original_error( + self, context, data_processor, monkeypatch + ): + import pandas + + self._use_schema(context, monkeypatch, {"r": "LARGE_BINARY"}) + with pytest.raises(TypeError) as exc_info: + data_processor._set_output_tuple( + iter([{"r": pandas.DataFrame({"x": [1]})}]) + ) + message = str(exc_info.value) + assert "Unmatched type for field 'r'" in message + assert "column 'r'" in message + + def test_column_hint_mentions_the_source_udf_columns_setting( + self, context, data_processor, monkeypatch + ): + self._use_schema(context, monkeypatch, {"x": "INTEGER"}) + with pytest.raises(KeyError) as exc_info: + data_processor._set_output_tuple(iter([{"y": 1}])) + assert "'Columns'" in str(exc_info.value)
