diff --git a/amber/src/main/python/core/models/operator.py b/amber/src/main/python/core/models/operator.py index 4e65fbb2ab2..47ca66c1c7d 100644 --- a/amber/src/main/python/core/models/operator.py +++ b/amber/src/main/python/core/models/operator.py @@ -22,7 +22,7 @@ from collections import defaultdict from typing import Iterator, List, Mapping, Optional, Union, MutableMapping, Protocol -from . import Table, TableLike, Tuple, TupleLike, Batch, BatchLike +from . import Table, TableLike, Tuple, TupleLike, Batch, BatchLike, Schema from .state import State from .table import all_output_to_tuple @@ -92,6 +92,23 @@ def is_source(self) -> bool: def is_source(self, value: bool) -> None: self.__internal_is_source = value + __internal_input_schemas: Optional[MutableMapping[int, Schema]] = None + + @property + @overrides.final + def input_schemas(self) -> MutableMapping[int, Schema]: + """ + What each input port was declared to carry, keyed by port index and + written by the runtime before that port's data is handed over. + + The tuples themselves say this too, so an operator only needs to ask + when there are none: a port can finish having carried no rows at all, + and its schema is then the only record of what its columns were. + """ + if self.__internal_input_schemas is None: + self.__internal_input_schemas = {} + return self.__internal_input_schemas + def open(self) -> None: """ Open a context of the operator. Usually can be used for loading/initiating some @@ -276,7 +293,11 @@ def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike yield def on_finish(self, port: int) -> Iterator[Optional[TableLike]]: - table = Table(self.__table_data[port]) + rows = self.__table_data[port] + schema = self.input_schemas.get(port) + # A port that carried no rows has no tuples to read column names off, + # and a table of no columns fails every operator that names one. + table = Table(rows) if rows or schema is None else Table.empty_of(schema) yield from self.process_table(table, port) @abstractmethod diff --git a/amber/src/main/python/core/models/table.py b/amber/src/main/python/core/models/table.py index 4716e0eba06..708e0580f86 100644 --- a/amber/src/main/python/core/models/table.py +++ b/amber/src/main/python/core/models/table.py @@ -16,6 +16,7 @@ # under the License. import pandas +import pyarrow as pa from pampy import match from typing import Iterator, TypeVar, List @@ -25,6 +26,22 @@ class Table(pandas.DataFrame): + @staticmethod + def empty_of(schema) -> "Table": + """ + The declared columns with no rows under them. + + ``from_tuple_likes`` reads the column names off the tuples it is given, + so with none to read it produces a frame of no columns at all, and an + operator naming any of its own columns raises KeyError. A port that + carried no rows still has a schema, and this is what it looks like as a + table. Building it through Arrow gives each column the dtype it would + have had with rows in it. + """ + return Table( + pa.Table.from_pylist([], schema=schema.as_arrow_schema()).to_pandas() + ) + @staticmethod def from_table(table): return table diff --git a/amber/src/main/python/core/runnables/data_processor.py b/amber/src/main/python/core/runnables/data_processor.py index 22e7058f27d..698e6b6da3e 100644 --- a/amber/src/main/python/core/runnables/data_processor.py +++ b/amber/src/main/python/core/runnables/data_processor.py @@ -74,8 +74,22 @@ def process_internal_marker(self, internal_marker: InternalMarker) -> None: # 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._declare_input_schema(executor, port_id) self._set_output_tuple(executor.on_finish(port_id)) + def _declare_input_schema(self, executor, port_id: int) -> None: + """ + Tell the executor what the finishing port was declared to carry, so an + operator handed no rows can still say what its columns were. A source + has no input port to ask, so it is left alone. + """ + port_identity = self._context.tuple_processing_manager.current_input_port_id + if port_identity is None: + return + executor.input_schemas[port_id] = self._context.input_manager.get_port( + port_identity + ).get_schema() + def process_state(self, state: State) -> None: """ Process an input marker by invoking appropriate state diff --git a/amber/src/test/python/core/models/test_operator.py b/amber/src/test/python/core/models/test_operator.py index bb18853350e..956aca8dba6 100644 --- a/amber/src/test/python/core/models/test_operator.py +++ b/amber/src/test/python/core/models/test_operator.py @@ -22,6 +22,7 @@ from core.models import ( BatchOperator, + Schema, SourceOperator, State, Table, @@ -386,6 +387,24 @@ def test_on_finish_with_no_buffered_tuples_yields_empty_table(self): assert len(op.received_tables) == 1 assert list(op.received_tables[0].as_tuples()) == [] + def test_on_finish_with_no_rows_keeps_the_declared_columns(self): + # The column names are read off the tuples, so a port that carried none + # left the operator a frame of no columns and every operator naming one + # of its own raised KeyError. The port's schema is the only record left. + op = _ConcreteTable() + op.input_schemas[0] = Schema(raw_schema={"x": "INTEGER", "y": "STRING"}) + + list(op.on_finish(port=0)) + + table = op.received_tables[0] + # Still a Table, so an operator reading it with as_tuples() keeps working. + assert isinstance(table, Table) + assert list(table.as_tuples()) == [] + assert list(table.columns) == ["x", "y"] + assert table.empty + # The dtype each column would have had with rows under it. + assert table["x"].dtype == "int32" + def test_buffers_are_keyed_by_port(self): # Each input port has its own tuple buffer; on_finish for one port # must not surface tuples written through a different port. diff --git a/amber/src/test/python/core/runnables/test_main_loop.py b/amber/src/test/python/core/runnables/test_main_loop.py index c5a566ada9c..2d59bf3f7ee 100644 --- a/amber/src/test/python/core/runnables/test_main_loop.py +++ b/amber/src/test/python/core/runnables/test_main_loop.py @@ -215,6 +215,10 @@ def state_processing_executor(self): # from `produce_state_on_finish` so EndChannel handling can be # observed. class StateProcessingExecutor: + # What Operator gives a real executor, for DataProcessor to record + # the finishing port's declared schema in. + input_schemas: dict = {} + @staticmethod def process_tuple(tuple_, port): yield tuple_