| # |
| # 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. |
| # |
| |
| """SDK harness for executing Python Fns via the Fn API.""" |
| |
| # pytype: skip-file |
| |
| import base64 |
| import bisect |
| import collections |
| import copy |
| import json |
| import logging |
| import random |
| import threading |
| from dataclasses import dataclass |
| from dataclasses import field |
| from typing import TYPE_CHECKING |
| from typing import Any |
| from typing import Callable |
| from typing import Container |
| from typing import DefaultDict |
| from typing import Dict |
| from typing import FrozenSet |
| from typing import Iterable |
| from typing import Iterator |
| from typing import List |
| from typing import Mapping |
| from typing import Optional |
| from typing import Set |
| from typing import Tuple |
| from typing import Type |
| from typing import TypeVar |
| from typing import Union |
| from typing import cast |
| |
| from google.protobuf import duration_pb2 |
| from google.protobuf import timestamp_pb2 |
| |
| import apache_beam as beam |
| from apache_beam import coders |
| from apache_beam.coders import WindowedValueCoder |
| from apache_beam.coders import coder_impl |
| from apache_beam.internal import pickler |
| from apache_beam.io import iobase |
| from apache_beam.metrics import monitoring_infos |
| from apache_beam.portability import common_urns |
| from apache_beam.portability import python_urns |
| from apache_beam.portability.api import beam_fn_api_pb2 |
| from apache_beam.portability.api import beam_runner_api_pb2 |
| from apache_beam.runners import common |
| from apache_beam.runners import pipeline_context |
| from apache_beam.runners.worker import data_sampler |
| from apache_beam.runners.worker import operation_specs |
| from apache_beam.runners.worker import operations |
| from apache_beam.runners.worker import statesampler |
| from apache_beam.transforms import TimeDomain |
| from apache_beam.transforms import core |
| from apache_beam.transforms import environments |
| from apache_beam.transforms import sideinputs |
| from apache_beam.transforms import userstate |
| from apache_beam.transforms import window |
| from apache_beam.utils import counters |
| from apache_beam.utils import proto_utils |
| from apache_beam.utils import timestamp |
| from apache_beam.utils.windowed_value import WindowedValue |
| |
| if TYPE_CHECKING: |
| from google.protobuf import message # pylint: disable=ungrouped-imports |
| from apache_beam import pvalue |
| from apache_beam.portability.api import metrics_pb2 |
| from apache_beam.runners.sdf_utils import SplitResultPrimary |
| from apache_beam.runners.sdf_utils import SplitResultResidual |
| from apache_beam.runners.worker import data_plane |
| from apache_beam.runners.worker import sdk_worker |
| from apache_beam.transforms.core import Windowing |
| from apache_beam.transforms.window import BoundedWindow |
| from apache_beam.utils import windowed_value |
| |
| T = TypeVar('T') |
| ConstructorFn = Callable[[ |
| 'BeamTransformFactory', |
| Any, |
| beam_runner_api_pb2.PTransform, |
| Union['message.Message', bytes], |
| Dict[str, List[operations.Operation]] |
| ], |
| operations.Operation] |
| OperationT = TypeVar('OperationT', bound=operations.Operation) |
| FnApiUserRuntimeStateTypes = Union['ReadModifyWriteRuntimeState', |
| 'CombiningValueRuntimeState', |
| 'SynchronousSetRuntimeState', |
| 'SynchronousBagRuntimeState'] |
| |
| DATA_INPUT_URN = 'beam:runner:source:v1' |
| DATA_OUTPUT_URN = 'beam:runner:sink:v1' |
| SYNTHETIC_DATA_SAMPLING_URN = 'beam:internal:sampling:v1' |
| IDENTITY_DOFN_URN = 'beam:dofn:identity:0.1' |
| # TODO(vikasrk): Fix this once runner sends appropriate common_urns. |
| OLD_DATAFLOW_RUNNER_HARNESS_PARDO_URN = 'beam:dofn:javasdk:0.1' |
| OLD_DATAFLOW_RUNNER_HARNESS_READ_URN = 'beam:source:java:0.1' |
| URNS_NEEDING_PCOLLECTIONS = set([ |
| monitoring_infos.ELEMENT_COUNT_URN, monitoring_infos.SAMPLED_BYTE_SIZE_URN |
| ]) |
| |
| _LOGGER = logging.getLogger(__name__) |
| |
| |
| class RunnerIOOperation(operations.Operation): |
| """Common baseclass for runner harness IO operations.""" |
| |
| def __init__(self, |
| name_context, # type: common.NameContext |
| step_name, # type: Any |
| consumers, # type: Mapping[Any, Iterable[operations.Operation]] |
| counter_factory, # type: counters.CounterFactory |
| state_sampler, # type: statesampler.StateSampler |
| windowed_coder, # type: coders.Coder |
| transform_id, # type: str |
| data_channel # type: data_plane.DataChannel |
| ): |
| # type: (...) -> None |
| super().__init__(name_context, None, counter_factory, state_sampler) |
| self.windowed_coder = windowed_coder |
| self.windowed_coder_impl = windowed_coder.get_impl() |
| # transform_id represents the consumer for the bytes in the data plane for a |
| # DataInputOperation or a producer of these bytes for a DataOutputOperation. |
| self.transform_id = transform_id |
| self.data_channel = data_channel |
| for _, consumer_ops in consumers.items(): |
| for consumer in consumer_ops: |
| self.add_receiver(consumer, 0) |
| |
| |
| class DataOutputOperation(RunnerIOOperation): |
| """A sink-like operation that gathers outputs to be sent back to the runner. |
| """ |
| def set_output_stream(self, output_stream): |
| # type: (data_plane.ClosableOutputStream) -> None |
| self.output_stream = output_stream |
| |
| def process(self, windowed_value): |
| # type: (windowed_value.WindowedValue) -> None |
| self.windowed_coder_impl.encode_to_stream( |
| windowed_value, self.output_stream, True) |
| self.output_stream.maybe_flush() |
| |
| def finish(self): |
| # type: () -> None |
| super().finish() |
| self.output_stream.close() |
| |
| |
| class DataInputOperation(RunnerIOOperation): |
| """A source-like operation that gathers input from the runner.""" |
| |
| def __init__(self, |
| operation_name, # type: common.NameContext |
| step_name, |
| consumers, # type: Mapping[Any, List[operations.Operation]] |
| counter_factory, # type: counters.CounterFactory |
| state_sampler, # type: statesampler.StateSampler |
| windowed_coder, # type: coders.Coder |
| transform_id, |
| data_channel # type: data_plane.GrpcClientDataChannel |
| ): |
| # type: (...) -> None |
| super().__init__( |
| operation_name, |
| step_name, |
| consumers, |
| counter_factory, |
| state_sampler, |
| windowed_coder, |
| transform_id=transform_id, |
| data_channel=data_channel) |
| |
| self.consumer = next(iter(consumers.values())) |
| self.splitting_lock = threading.Lock() |
| self.index = -1 |
| self.stop = float('inf') |
| self.started = False |
| |
| def setup(self, data_sampler=None): |
| super().setup(data_sampler) |
| # We must do this manually as we don't have a spec or spec.output_coders. |
| self.receivers = [ |
| operations.ConsumerSet.create( |
| counter_factory=self.counter_factory, |
| step_name=self.name_context.step_name, |
| output_index=0, |
| consumers=self.consumer, |
| coder=self.windowed_coder, |
| producer_type_hints=self._get_runtime_performance_hints(), |
| producer_batch_converter=self.get_output_batch_converter()) |
| ] |
| |
| def start(self): |
| # type: () -> None |
| super().start() |
| with self.splitting_lock: |
| self.started = True |
| |
| def process(self, windowed_value): |
| # type: (windowed_value.WindowedValue) -> None |
| self.output(windowed_value) |
| |
| def process_encoded(self, encoded_windowed_values): |
| # type: (bytes) -> None |
| input_stream = coder_impl.create_InputStream(encoded_windowed_values) |
| while input_stream.size() > 0: |
| with self.splitting_lock: |
| if self.index == self.stop - 1: |
| return |
| self.index += 1 |
| try: |
| decoded_value = self.windowed_coder_impl.decode_from_stream( |
| input_stream, True) |
| except Exception as exn: |
| raise ValueError( |
| "Error decoding input stream with coder " + |
| str(self.windowed_coder)) from exn |
| self.output(decoded_value) |
| |
| def monitoring_infos(self, transform_id, tag_to_pcollection_id): |
| # type: (str, Dict[str, str]) -> Dict[FrozenSet, metrics_pb2.MonitoringInfo] |
| all_monitoring_infos = super().monitoring_infos( |
| transform_id, tag_to_pcollection_id) |
| read_progress_info = monitoring_infos.int64_counter( |
| monitoring_infos.DATA_CHANNEL_READ_INDEX, |
| self.index, |
| ptransform=transform_id) |
| all_monitoring_infos[monitoring_infos.to_key( |
| read_progress_info)] = read_progress_info |
| return all_monitoring_infos |
| |
| # TODO(https://github.com/apache/beam/issues/19737): typing not compatible |
| # with super type |
| def try_split( # type: ignore[override] |
| self, fraction_of_remainder, total_buffer_size, allowed_split_points): |
| # type: (...) -> Optional[Tuple[int, Iterable[operations.SdfSplitResultsPrimary], Iterable[operations.SdfSplitResultsResidual], int]] |
| with self.splitting_lock: |
| if not self.started: |
| return None |
| if self.index == -1: |
| # We are "finished" with the (non-existent) previous element. |
| current_element_progress = 1.0 |
| else: |
| current_element_progress_object = ( |
| self.receivers[0].current_element_progress()) |
| if current_element_progress_object is None: |
| current_element_progress = 0.5 |
| else: |
| current_element_progress = ( |
| current_element_progress_object.fraction_completed) |
| # Now figure out where to split. |
| split = self._compute_split( |
| self.index, |
| current_element_progress, |
| self.stop, |
| fraction_of_remainder, |
| total_buffer_size, |
| allowed_split_points, |
| self.receivers[0].try_split) |
| if split: |
| self.stop = split[-1] |
| return split |
| |
| @staticmethod |
| def _compute_split( |
| index, |
| current_element_progress, |
| stop, |
| fraction_of_remainder, |
| total_buffer_size, |
| allowed_split_points=(), |
| try_split=lambda fraction: None): |
| def is_valid_split_point(index): |
| return not allowed_split_points or index in allowed_split_points |
| |
| if total_buffer_size < index + 1: |
| total_buffer_size = index + 1 |
| elif total_buffer_size > stop: |
| total_buffer_size = stop |
| # The units here (except for keep_of_element_remainder) are all in |
| # terms of number of (possibly fractional) elements. |
| remainder = total_buffer_size - index - current_element_progress |
| keep = remainder * fraction_of_remainder |
| if current_element_progress < 1: |
| keep_of_element_remainder = keep / (1 - current_element_progress) |
| # If it's less than what's left of the current element, |
| # try splitting at the current element. |
| if (keep_of_element_remainder < 1 and is_valid_split_point(index) and |
| is_valid_split_point(index + 1)): |
| split = try_split( |
| keep_of_element_remainder |
| ) # type: Optional[Tuple[Iterable[operations.SdfSplitResultsPrimary], Iterable[operations.SdfSplitResultsResidual]]] |
| if split: |
| element_primaries, element_residuals = split |
| return index - 1, element_primaries, element_residuals, index + 1 |
| # Otherwise, split at the closest element boundary. |
| # pylint: disable=bad-option-value |
| stop_index = index + max(1, int(round(current_element_progress + keep))) |
| if allowed_split_points and stop_index not in allowed_split_points: |
| # Choose the closest allowed split point. |
| allowed_split_points = sorted(allowed_split_points) |
| closest = bisect.bisect(allowed_split_points, stop_index) |
| if closest == 0: |
| stop_index = allowed_split_points[0] |
| elif closest == len(allowed_split_points): |
| stop_index = allowed_split_points[-1] |
| else: |
| prev = allowed_split_points[closest - 1] |
| next = allowed_split_points[closest] |
| if index < prev and stop_index - prev < next - stop_index: |
| stop_index = prev |
| else: |
| stop_index = next |
| if index < stop_index < stop: |
| return stop_index - 1, [], [], stop_index |
| else: |
| return None |
| |
| def finish(self): |
| # type: () -> None |
| super().finish() |
| with self.splitting_lock: |
| self.index += 1 |
| self.started = False |
| |
| def reset(self): |
| # type: () -> None |
| with self.splitting_lock: |
| self.index = -1 |
| self.stop = float('inf') |
| super().reset() |
| |
| |
| class _StateBackedIterable(object): |
| def __init__(self, |
| state_handler, # type: sdk_worker.CachingStateHandler |
| state_key, # type: beam_fn_api_pb2.StateKey |
| coder_or_impl, # type: Union[coders.Coder, coder_impl.CoderImpl] |
| ): |
| # type: (...) -> None |
| self._state_handler = state_handler |
| self._state_key = state_key |
| if isinstance(coder_or_impl, coders.Coder): |
| self._coder_impl = coder_or_impl.get_impl() |
| else: |
| self._coder_impl = coder_or_impl |
| |
| def __iter__(self): |
| # type: () -> Iterator[Any] |
| return iter( |
| self._state_handler.blocking_get(self._state_key, self._coder_impl)) |
| |
| def __reduce__(self): |
| return list, (list(self), ) |
| |
| |
| coder_impl.FastPrimitivesCoderImpl.register_iterable_like_type( |
| _StateBackedIterable) |
| |
| |
| class StateBackedSideInputMap(object): |
| |
| _BULK_READ_LIMIT = 100 |
| _BULK_READ_FULLY = "fully" |
| _BULK_READ_PARTIALLY = "partially" |
| |
| def __init__(self, |
| state_handler, # type: sdk_worker.CachingStateHandler |
| transform_id, # type: str |
| tag, # type: Optional[str] |
| side_input_data, # type: pvalue.SideInputData |
| coder, # type: WindowedValueCoder |
| use_bulk_read = False, # type: bool |
| ): |
| # type: (...) -> None |
| self._state_handler = state_handler |
| self._transform_id = transform_id |
| self._tag = tag |
| self._side_input_data = side_input_data |
| self._element_coder = coder.wrapped_value_coder |
| self._target_window_coder = coder.window_coder |
| # TODO(robertwb): Limit the cache size. |
| self._cache = {} # type: Dict[BoundedWindow, Any] |
| self._use_bulk_read = use_bulk_read |
| |
| def __getitem__(self, window): |
| target_window = self._side_input_data.window_mapping_fn(window) |
| if target_window not in self._cache: |
| state_handler = self._state_handler |
| access_pattern = self._side_input_data.access_pattern |
| |
| if access_pattern == common_urns.side_inputs.ITERABLE.urn: |
| state_key = beam_fn_api_pb2.StateKey( |
| iterable_side_input=beam_fn_api_pb2.StateKey.IterableSideInput( |
| transform_id=self._transform_id, |
| side_input_id=self._tag, |
| window=self._target_window_coder.encode(target_window))) |
| raw_view = _StateBackedIterable( |
| state_handler, state_key, self._element_coder) |
| |
| elif access_pattern == common_urns.side_inputs.MULTIMAP.urn: |
| state_key = beam_fn_api_pb2.StateKey( |
| multimap_side_input=beam_fn_api_pb2.StateKey.MultimapSideInput( |
| transform_id=self._transform_id, |
| side_input_id=self._tag, |
| window=self._target_window_coder.encode(target_window), |
| key=b'')) |
| kv_iter_state_key = beam_fn_api_pb2.StateKey( |
| multimap_keys_values_side_input=beam_fn_api_pb2.StateKey. |
| MultimapKeysValuesSideInput( |
| transform_id=self._transform_id, |
| side_input_id=self._tag, |
| window=self._target_window_coder.encode(target_window))) |
| cache = {} |
| key_coder = self._element_coder.key_coder() |
| key_coder_impl = key_coder.get_impl() |
| value_coder = self._element_coder.value_coder() |
| use_bulk_read = self._use_bulk_read |
| |
| class MultiMap(object): |
| _bulk_read = None |
| _lock = threading.Lock() |
| |
| def __getitem__(self, key): |
| if use_bulk_read: |
| if self._bulk_read is None: |
| with self._lock: |
| if self._bulk_read is None: |
| try: |
| # Attempt to bulk read the key-values over the iterable |
| # protocol which, if supported, can be much more efficient |
| # than point lookups if it fits into memory. |
| for ix, (k, vs) in enumerate(_StateBackedIterable( |
| state_handler, |
| kv_iter_state_key, |
| coders.TupleCoder( |
| (key_coder, coders.IterableCoder(value_coder))))): |
| cache[k] = vs |
| if ix > StateBackedSideInputMap._BULK_READ_LIMIT: |
| self._bulk_read = ( |
| StateBackedSideInputMap._BULK_READ_PARTIALLY) |
| break |
| else: |
| # We reached the end of the iteration without breaking. |
| self._bulk_read = ( |
| StateBackedSideInputMap._BULK_READ_FULLY) |
| except Exception: |
| _LOGGER.error( |
| "Iterable access of map side inputs unsupported.", |
| exc_info=True) |
| self._bulk_read = ( |
| StateBackedSideInputMap._BULK_READ_PARTIALLY) |
| |
| if (self._bulk_read == StateBackedSideInputMap._BULK_READ_FULLY): |
| return cache.get(key, []) |
| |
| if key not in cache: |
| keyed_state_key = beam_fn_api_pb2.StateKey() |
| keyed_state_key.CopyFrom(state_key) |
| keyed_state_key.multimap_side_input.key = ( |
| key_coder_impl.encode_nested(key)) |
| cache[key] = _StateBackedIterable( |
| state_handler, keyed_state_key, value_coder) |
| |
| return cache[key] |
| |
| def __reduce__(self): |
| # TODO(robertwb): Figure out how to support this. |
| raise TypeError(common_urns.side_inputs.MULTIMAP.urn) |
| |
| raw_view = MultiMap() |
| |
| else: |
| raise ValueError("Unknown access pattern: '%s'" % access_pattern) |
| |
| self._cache[target_window] = self._side_input_data.view_fn(raw_view) |
| return self._cache[target_window] |
| |
| def is_globally_windowed(self): |
| # type: () -> bool |
| return ( |
| self._side_input_data.window_mapping_fn == |
| sideinputs._global_window_mapping_fn) |
| |
| def reset(self): |
| # type: () -> None |
| # TODO(BEAM-5428): Cross-bundle caching respecting cache tokens. |
| self._cache = {} |
| |
| |
| class ReadModifyWriteRuntimeState(userstate.ReadModifyWriteRuntimeState): |
| def __init__(self, underlying_bag_state): |
| self._underlying_bag_state = underlying_bag_state |
| |
| def read(self): # type: () -> Any |
| values = list(self._underlying_bag_state.read()) |
| if not values: |
| return None |
| return values[0] |
| |
| def write(self, value): # type: (Any) -> None |
| self.clear() |
| self._underlying_bag_state.add(value) |
| |
| def clear(self): # type: () -> None |
| self._underlying_bag_state.clear() |
| |
| def commit(self): # type: () -> None |
| self._underlying_bag_state.commit() |
| |
| |
| class CombiningValueRuntimeState(userstate.CombiningValueRuntimeState): |
| def __init__(self, underlying_bag_state, combinefn): |
| # type: (userstate.AccumulatingRuntimeState, core.CombineFn) -> None |
| self._combinefn = combinefn |
| self._combinefn.setup() |
| self._underlying_bag_state = underlying_bag_state |
| self._finalized = False |
| |
| def _read_accumulator(self, rewrite=True): |
| merged_accumulator = self._combinefn.merge_accumulators( |
| self._underlying_bag_state.read()) |
| if rewrite: |
| self._underlying_bag_state.clear() |
| self._underlying_bag_state.add(merged_accumulator) |
| return merged_accumulator |
| |
| def read(self): |
| # type: () -> Iterable[Any] |
| return self._combinefn.extract_output(self._read_accumulator()) |
| |
| def add(self, value): |
| # type: (Any) -> None |
| # Prefer blind writes, but don't let them grow unboundedly. |
| # This should be tuned to be much lower, but for now exercise |
| # both paths well. |
| if random.random() < 0.5: |
| accumulator = self._read_accumulator(False) |
| self._underlying_bag_state.clear() |
| else: |
| accumulator = self._combinefn.create_accumulator() |
| self._underlying_bag_state.add( |
| self._combinefn.add_input(accumulator, value)) |
| |
| def clear(self): |
| # type: () -> None |
| self._underlying_bag_state.clear() |
| |
| def commit(self): |
| self._underlying_bag_state.commit() |
| |
| def finalize(self): |
| if not self._finalized: |
| self._combinefn.teardown() |
| self._finalized = True |
| |
| |
| class _ConcatIterable(object): |
| """An iterable that is the concatination of two iterables. |
| |
| Unlike itertools.chain, this allows reiteration. |
| """ |
| def __init__(self, first, second): |
| # type: (Iterable[Any], Iterable[Any]) -> None |
| self.first = first |
| self.second = second |
| |
| def __iter__(self): |
| # type: () -> Iterator[Any] |
| for elem in self.first: |
| yield elem |
| for elem in self.second: |
| yield elem |
| |
| |
| coder_impl.FastPrimitivesCoderImpl.register_iterable_like_type(_ConcatIterable) |
| |
| |
| class SynchronousBagRuntimeState(userstate.BagRuntimeState): |
| |
| def __init__(self, |
| state_handler, # type: sdk_worker.CachingStateHandler |
| state_key, # type: beam_fn_api_pb2.StateKey |
| value_coder # type: coders.Coder |
| ): |
| # type: (...) -> None |
| self._state_handler = state_handler |
| self._state_key = state_key |
| self._value_coder = value_coder |
| self._cleared = False |
| self._added_elements = [] # type: List[Any] |
| |
| def read(self): |
| # type: () -> Iterable[Any] |
| return _ConcatIterable([] if self._cleared else cast( |
| 'Iterable[Any]', |
| _StateBackedIterable( |
| self._state_handler, self._state_key, self._value_coder)), |
| self._added_elements) |
| |
| def add(self, value): |
| # type: (Any) -> None |
| self._added_elements.append(value) |
| |
| def clear(self): |
| # type: () -> None |
| self._cleared = True |
| self._added_elements = [] |
| |
| def commit(self): |
| # type: () -> None |
| to_await = None |
| if self._cleared: |
| to_await = self._state_handler.clear(self._state_key) |
| if self._added_elements: |
| to_await = self._state_handler.extend( |
| self._state_key, self._value_coder.get_impl(), self._added_elements) |
| if to_await: |
| # To commit, we need to wait on the last state request future to complete. |
| to_await.get() |
| |
| |
| class SynchronousSetRuntimeState(userstate.SetRuntimeState): |
| |
| def __init__(self, |
| state_handler, # type: sdk_worker.CachingStateHandler |
| state_key, # type: beam_fn_api_pb2.StateKey |
| value_coder # type: coders.Coder |
| ): |
| # type: (...) -> None |
| self._state_handler = state_handler |
| self._state_key = state_key |
| self._value_coder = value_coder |
| self._cleared = False |
| self._added_elements = set() # type: Set[Any] |
| |
| def _compact_data(self, rewrite=True): |
| accumulator = set( |
| _ConcatIterable( |
| set() if self._cleared else _StateBackedIterable( |
| self._state_handler, self._state_key, self._value_coder), |
| self._added_elements)) |
| |
| if rewrite and accumulator: |
| self._state_handler.clear(self._state_key) |
| self._state_handler.extend( |
| self._state_key, self._value_coder.get_impl(), accumulator) |
| |
| # Since everthing is already committed so we can safely reinitialize |
| # added_elements here. |
| self._added_elements = set() |
| |
| return accumulator |
| |
| def read(self): |
| # type: () -> Set[Any] |
| return self._compact_data(rewrite=False) |
| |
| def add(self, value): |
| # type: (Any) -> None |
| if self._cleared: |
| # This is a good time explicitly clear. |
| self._state_handler.clear(self._state_key) |
| self._cleared = False |
| |
| self._added_elements.add(value) |
| if random.random() > 0.5: |
| self._compact_data() |
| |
| def clear(self): |
| # type: () -> None |
| self._cleared = True |
| self._added_elements = set() |
| |
| def commit(self): |
| # type: () -> None |
| to_await = None |
| if self._cleared: |
| to_await = self._state_handler.clear(self._state_key) |
| if self._added_elements: |
| to_await = self._state_handler.extend( |
| self._state_key, self._value_coder.get_impl(), self._added_elements) |
| if to_await: |
| # To commit, we need to wait on the last state request future to complete. |
| to_await.get() |
| |
| |
| class OutputTimer(userstate.BaseTimer): |
| def __init__(self, |
| key, |
| window, # type: BoundedWindow |
| timestamp, # type: timestamp.Timestamp |
| paneinfo, # type: windowed_value.PaneInfo |
| time_domain, # type: str |
| timer_family_id, # type: str |
| timer_coder_impl, # type: coder_impl.TimerCoderImpl |
| output_stream # type: data_plane.ClosableOutputStream |
| ): |
| self._key = key |
| self._window = window |
| self._input_timestamp = timestamp |
| self._paneinfo = paneinfo |
| self._time_domain = time_domain |
| self._timer_family_id = timer_family_id |
| self._output_stream = output_stream |
| self._timer_coder_impl = timer_coder_impl |
| |
| def set(self, ts: timestamp.TimestampTypes, dynamic_timer_tag='') -> None: |
| ts = timestamp.Timestamp.of(ts) |
| timer = userstate.Timer( |
| user_key=self._key, |
| dynamic_timer_tag=dynamic_timer_tag, |
| windows=(self._window, ), |
| clear_bit=False, |
| fire_timestamp=ts, |
| hold_timestamp=ts if TimeDomain.is_event_time(self._time_domain) else |
| self._input_timestamp, |
| paneinfo=self._paneinfo) |
| self._timer_coder_impl.encode_to_stream(timer, self._output_stream, True) |
| self._output_stream.maybe_flush() |
| |
| def clear(self, dynamic_timer_tag='') -> None: |
| timer = userstate.Timer( |
| user_key=self._key, |
| dynamic_timer_tag=dynamic_timer_tag, |
| windows=(self._window, ), |
| clear_bit=True, |
| fire_timestamp=None, |
| hold_timestamp=None, |
| paneinfo=None) |
| self._timer_coder_impl.encode_to_stream(timer, self._output_stream, True) |
| self._output_stream.maybe_flush() |
| |
| |
| class TimerInfo(object): |
| """A data class to store information related to a timer.""" |
| def __init__(self, timer_coder_impl, output_stream=None): |
| self.timer_coder_impl = timer_coder_impl |
| self.output_stream = output_stream |
| |
| |
| class FnApiUserStateContext(userstate.UserStateContext): |
| """Interface for state and timers from SDK to Fn API servicer of state..""" |
| |
| def __init__(self, |
| state_handler, # type: sdk_worker.CachingStateHandler |
| transform_id, # type: str |
| key_coder, # type: coders.Coder |
| window_coder, # type: coders.Coder |
| ): |
| # type: (...) -> None |
| |
| """Initialize a ``FnApiUserStateContext``. |
| |
| Args: |
| state_handler: A StateServicer object. |
| transform_id: The name of the PTransform that this context is associated. |
| key_coder: Coder for the key type. |
| window_coder: Coder for the window type. |
| """ |
| self._state_handler = state_handler |
| self._transform_id = transform_id |
| self._key_coder = key_coder |
| self._window_coder = window_coder |
| # A mapping of {timer_family_id: TimerInfo} |
| self._timers_info = {} # type: Dict[str, TimerInfo] |
| self._all_states = {} # type: Dict[tuple, FnApiUserRuntimeStateTypes] |
| |
| def add_timer_info(self, timer_family_id, timer_info): |
| # type: (str, TimerInfo) -> None |
| self._timers_info[timer_family_id] = timer_info |
| |
| def get_timer( |
| self, timer_spec: userstate.TimerSpec, key, window, timestamp, |
| pane) -> OutputTimer: |
| assert self._timers_info[timer_spec.name].output_stream is not None |
| timer_coder_impl = self._timers_info[timer_spec.name].timer_coder_impl |
| output_stream = self._timers_info[timer_spec.name].output_stream |
| return OutputTimer( |
| key, |
| window, |
| timestamp, |
| pane, |
| timer_spec.time_domain, |
| timer_spec.name, |
| timer_coder_impl, |
| output_stream) |
| |
| def get_state(self, *args): |
| # type: (*Any) -> FnApiUserRuntimeStateTypes |
| state_handle = self._all_states.get(args) |
| if state_handle is None: |
| state_handle = self._all_states[args] = self._create_state(*args) |
| return state_handle |
| |
| def _create_state(self, |
| state_spec, # type: userstate.StateSpec |
| key, |
| window # type: BoundedWindow |
| ): |
| # type: (...) -> FnApiUserRuntimeStateTypes |
| if isinstance(state_spec, |
| (userstate.BagStateSpec, |
| userstate.CombiningValueStateSpec, |
| userstate.ReadModifyWriteStateSpec)): |
| bag_state = SynchronousBagRuntimeState( |
| self._state_handler, |
| state_key=beam_fn_api_pb2.StateKey( |
| bag_user_state=beam_fn_api_pb2.StateKey.BagUserState( |
| transform_id=self._transform_id, |
| user_state_id=state_spec.name, |
| window=self._window_coder.encode(window), |
| # State keys are expected in nested encoding format |
| key=self._key_coder.encode_nested(key))), |
| value_coder=state_spec.coder) |
| if isinstance(state_spec, userstate.BagStateSpec): |
| return bag_state |
| elif isinstance(state_spec, userstate.ReadModifyWriteStateSpec): |
| return ReadModifyWriteRuntimeState(bag_state) |
| else: |
| return CombiningValueRuntimeState( |
| bag_state, copy.deepcopy(state_spec.combine_fn)) |
| elif isinstance(state_spec, userstate.SetStateSpec): |
| return SynchronousSetRuntimeState( |
| self._state_handler, |
| state_key=beam_fn_api_pb2.StateKey( |
| bag_user_state=beam_fn_api_pb2.StateKey.BagUserState( |
| transform_id=self._transform_id, |
| user_state_id=state_spec.name, |
| window=self._window_coder.encode(window), |
| # State keys are expected in nested encoding format |
| key=self._key_coder.encode_nested(key))), |
| value_coder=state_spec.coder) |
| else: |
| raise NotImplementedError(state_spec) |
| |
| def commit(self): |
| # type: () -> None |
| for state in self._all_states.values(): |
| state.commit() |
| |
| def reset(self): |
| # type: () -> None |
| for state in self._all_states.values(): |
| state.finalize() |
| self._all_states = {} |
| |
| |
| def memoize(func): |
| cache = {} |
| missing = object() |
| |
| def wrapper(*args): |
| result = cache.get(args, missing) |
| if result is missing: |
| result = cache[args] = func(*args) |
| return result |
| |
| return wrapper |
| |
| |
| def only_element(iterable): |
| # type: (Iterable[T]) -> T |
| element, = iterable |
| return element |
| |
| |
| def _environments_compatible(submission, runtime): |
| # type: (str, str) -> bool |
| if submission == runtime: |
| return True |
| if 'rc' in submission and runtime in submission: |
| # TODO(https://github.com/apache/beam/issues/28084): Loosen |
| # the check for RCs until RC containers install the matching version. |
| return True |
| return False |
| |
| |
| def _verify_descriptor_created_in_a_compatible_env(process_bundle_descriptor): |
| # type: (beam_fn_api_pb2.ProcessBundleDescriptor) -> None |
| |
| runtime_sdk = environments.sdk_base_version_capability() |
| for t in process_bundle_descriptor.transforms.values(): |
| env = process_bundle_descriptor.environments[t.environment_id] |
| for c in env.capabilities: |
| if (c.startswith(environments.SDK_VERSION_CAPABILITY_PREFIX) and |
| not _environments_compatible(c, runtime_sdk)): |
| raise RuntimeError( |
| "Pipeline construction environment and pipeline runtime " |
| "environment are not compatible. If you use a custom " |
| "container image, check that the Python interpreter minor version " |
| "and the Apache Beam version in your image match the versions " |
| "used at pipeline construction time. " |
| f"Submission environment: {c}. " |
| f"Runtime environment: {runtime_sdk}.") |
| |
| # TODO: Consider warning on mismatches in versions of installed packages. |
| |
| |
| class BundleProcessor(object): |
| """ A class for processing bundles of elements. """ |
| |
| def __init__(self, |
| runner_capabilities, # type: FrozenSet[str] |
| process_bundle_descriptor, # type: beam_fn_api_pb2.ProcessBundleDescriptor |
| state_handler, # type: sdk_worker.CachingStateHandler |
| data_channel_factory, # type: data_plane.DataChannelFactory |
| data_sampler=None, # type: Optional[data_sampler.DataSampler] |
| ): |
| # type: (...) -> None |
| |
| """Initialize a bundle processor. |
| |
| Args: |
| runner_capabilities (``FrozenSet[str]``): The set of capabilities of the |
| runner with which we will be interacting |
| process_bundle_descriptor (``beam_fn_api_pb2.ProcessBundleDescriptor``): |
| a description of the stage that this ``BundleProcessor``is to execute. |
| state_handler (CachingStateHandler). |
| data_channel_factory (``data_plane.DataChannelFactory``). |
| """ |
| self.runner_capabilities = runner_capabilities |
| self.process_bundle_descriptor = process_bundle_descriptor |
| self.state_handler = state_handler |
| self.data_channel_factory = data_channel_factory |
| self.data_sampler = data_sampler |
| self.current_instruction_id = None # type: Optional[str] |
| |
| _verify_descriptor_created_in_a_compatible_env(process_bundle_descriptor) |
| # There is no guarantee that the runner only set |
| # timer_api_service_descriptor when having timers. So this field cannot be |
| # used as an indicator of timers. |
| if self.process_bundle_descriptor.timer_api_service_descriptor.url: |
| self.timer_data_channel = ( |
| data_channel_factory.create_data_channel_from_url( |
| self.process_bundle_descriptor.timer_api_service_descriptor.url)) |
| else: |
| self.timer_data_channel = None |
| |
| # A mapping of |
| # {(transform_id, timer_family_id): TimerInfo} |
| # The mapping is empty when there is no timer_family_specs in the |
| # ProcessBundleDescriptor. |
| self.timers_info = {} # type: Dict[Tuple[str, str], TimerInfo] |
| |
| # TODO(robertwb): Figure out the correct prefix to use for output counters |
| # from StateSampler. |
| self.counter_factory = counters.CounterFactory() |
| self.state_sampler = statesampler.StateSampler( |
| 'fnapi-step-%s' % self.process_bundle_descriptor.id, |
| self.counter_factory) |
| |
| self.ops = self.create_execution_tree(self.process_bundle_descriptor) |
| for op in reversed(self.ops.values()): |
| op.setup(self.data_sampler) |
| self.splitting_lock = threading.Lock() |
| |
| def create_execution_tree( |
| self, |
| descriptor # type: beam_fn_api_pb2.ProcessBundleDescriptor |
| ): |
| # type: (...) -> collections.OrderedDict[str, operations.DoOperation] |
| transform_factory = BeamTransformFactory( |
| self.runner_capabilities, |
| descriptor, |
| self.data_channel_factory, |
| self.counter_factory, |
| self.state_sampler, |
| self.state_handler, |
| self.data_sampler, |
| ) |
| |
| self.timers_info = transform_factory.extract_timers_info() |
| |
| def is_side_input(transform_proto, tag): |
| if transform_proto.spec.urn == common_urns.primitives.PAR_DO.urn: |
| return tag in proto_utils.parse_Bytes( |
| transform_proto.spec.payload, |
| beam_runner_api_pb2.ParDoPayload).side_inputs |
| |
| pcoll_consumers = collections.defaultdict( |
| list) # type: DefaultDict[str, List[str]] |
| for transform_id, transform_proto in descriptor.transforms.items(): |
| for tag, pcoll_id in transform_proto.inputs.items(): |
| if not is_side_input(transform_proto, tag): |
| pcoll_consumers[pcoll_id].append(transform_id) |
| |
| @memoize |
| def get_operation(transform_id): |
| # type: (str) -> operations.Operation |
| transform_consumers = { |
| tag: [get_operation(op) for op in pcoll_consumers[pcoll_id]] |
| for tag, |
| pcoll_id in descriptor.transforms[transform_id].outputs.items() |
| } |
| |
| # Initialize transform-specific state in the Data Sampler. |
| if self.data_sampler: |
| self.data_sampler.initialize_samplers( |
| transform_id, descriptor, transform_factory.get_coder) |
| |
| return transform_factory.create_operation( |
| transform_id, transform_consumers) |
| |
| # Operations must be started (hence returned) in order. |
| @memoize |
| def topological_height(transform_id): |
| # type: (str) -> int |
| return 1 + max([0] + [ |
| topological_height(consumer) |
| for pcoll in descriptor.transforms[transform_id].outputs.values() |
| for consumer in pcoll_consumers[pcoll] |
| ]) |
| |
| return collections.OrderedDict([( |
| transform_id, |
| cast(operations.DoOperation, |
| get_operation(transform_id))) for transform_id in sorted( |
| descriptor.transforms, key=topological_height, reverse=True)]) |
| |
| def reset(self): |
| # type: () -> None |
| self.counter_factory.reset() |
| self.state_sampler.reset() |
| # Side input caches. |
| for op in self.ops.values(): |
| op.reset() |
| |
| def process_bundle(self, instruction_id): |
| # type: (str) -> Tuple[List[beam_fn_api_pb2.DelayedBundleApplication], bool] |
| |
| expected_input_ops = [] # type: List[DataInputOperation] |
| |
| for op in self.ops.values(): |
| if isinstance(op, DataOutputOperation): |
| # TODO(robertwb): Is there a better way to pass the instruction id to |
| # the operation? |
| op.set_output_stream( |
| op.data_channel.output_stream(instruction_id, op.transform_id)) |
| elif isinstance(op, DataInputOperation): |
| # We must wait until we receive "end of stream" for each of these ops. |
| expected_input_ops.append(op) |
| |
| try: |
| execution_context = ExecutionContext(instruction_id=instruction_id) |
| self.current_instruction_id = instruction_id |
| self.state_sampler.start() |
| # Start all operations. |
| for op in reversed(self.ops.values()): |
| _LOGGER.debug('start %s', op) |
| op.execution_context = execution_context |
| op.start() |
| |
| # Each data_channel is mapped to a list of expected inputs which includes |
| # both data input and timer input. The data input is identied by |
| # transform_id. The data input is identified by |
| # (transform_id, timer_family_id). |
| data_channels = collections.defaultdict( |
| list |
| ) # type: DefaultDict[data_plane.DataChannel, List[Union[str, Tuple[str, str]]]] |
| |
| # Add expected data inputs for each data channel. |
| input_op_by_transform_id = {} |
| for input_op in expected_input_ops: |
| data_channels[input_op.data_channel].append(input_op.transform_id) |
| input_op_by_transform_id[input_op.transform_id] = input_op |
| |
| # Update timer_data channel with expected timer inputs. |
| if self.timer_data_channel: |
| data_channels[self.timer_data_channel].extend( |
| list(self.timers_info.keys())) |
| |
| # Set up timer output stream for DoOperation. |
| for ((transform_id, timer_family_id), |
| timer_info) in self.timers_info.items(): |
| output_stream = self.timer_data_channel.output_timer_stream( |
| instruction_id, transform_id, timer_family_id) |
| timer_info.output_stream = output_stream |
| self.ops[transform_id].add_timer_info(timer_family_id, timer_info) |
| |
| # Process data and timer inputs |
| for data_channel, expected_inputs in data_channels.items(): |
| for element in data_channel.input_elements(instruction_id, |
| expected_inputs): |
| if isinstance(element, beam_fn_api_pb2.Elements.Timers): |
| timer_coder_impl = ( |
| self.timers_info[( |
| element.transform_id, |
| element.timer_family_id)].timer_coder_impl) |
| for timer_data in timer_coder_impl.decode_all(element.timers): |
| self.ops[element.transform_id].process_timer( |
| element.timer_family_id, timer_data) |
| elif isinstance(element, beam_fn_api_pb2.Elements.Data): |
| input_op_by_transform_id[element.transform_id].process_encoded( |
| element.data) |
| |
| # Finish all operations. |
| for op in self.ops.values(): |
| _LOGGER.debug('finish %s', op) |
| op.finish() |
| |
| # Close every timer output stream |
| for timer_info in self.timers_info.values(): |
| assert timer_info.output_stream is not None |
| timer_info.output_stream.close() |
| |
| return ([ |
| self.delayed_bundle_application(op, residual) for op, |
| residual in execution_context.delayed_applications |
| ], |
| self.requires_finalization()) |
| |
| finally: |
| # Ensure any in-flight split attempts complete. |
| with self.splitting_lock: |
| self.current_instruction_id = None |
| self.state_sampler.stop_if_still_running() |
| |
| def finalize_bundle(self): |
| # type: () -> beam_fn_api_pb2.FinalizeBundleResponse |
| for op in self.ops.values(): |
| op.finalize_bundle() |
| return beam_fn_api_pb2.FinalizeBundleResponse() |
| |
| def requires_finalization(self): |
| # type: () -> bool |
| return any(op.needs_finalization() for op in self.ops.values()) |
| |
| def try_split(self, bundle_split_request): |
| # type: (beam_fn_api_pb2.ProcessBundleSplitRequest) -> beam_fn_api_pb2.ProcessBundleSplitResponse |
| split_response = beam_fn_api_pb2.ProcessBundleSplitResponse() |
| with self.splitting_lock: |
| if bundle_split_request.instruction_id != self.current_instruction_id: |
| # This may be a delayed split for a former bundle, see BEAM-12475. |
| return split_response |
| |
| for op in self.ops.values(): |
| if isinstance(op, DataInputOperation): |
| desired_split = bundle_split_request.desired_splits.get( |
| op.transform_id) |
| if desired_split: |
| split = op.try_split( |
| desired_split.fraction_of_remainder, |
| desired_split.estimated_input_elements, |
| desired_split.allowed_split_points) |
| if split: |
| ( |
| primary_end, |
| element_primaries, |
| element_residuals, |
| residual_start, |
| ) = split |
| for element_primary in element_primaries: |
| split_response.primary_roots.add().CopyFrom( |
| self.bundle_application(*element_primary)) |
| for element_residual in element_residuals: |
| split_response.residual_roots.add().CopyFrom( |
| self.delayed_bundle_application(*element_residual)) |
| split_response.channel_splits.extend([ |
| beam_fn_api_pb2.ProcessBundleSplitResponse.ChannelSplit( |
| transform_id=op.transform_id, |
| last_primary_element=primary_end, |
| first_residual_element=residual_start) |
| ]) |
| |
| return split_response |
| |
| def delayed_bundle_application(self, |
| op, # type: operations.DoOperation |
| deferred_remainder # type: SplitResultResidual |
| ): |
| # type: (...) -> beam_fn_api_pb2.DelayedBundleApplication |
| assert op.input_info is not None |
| # TODO(SDF): For non-root nodes, need main_input_coder + residual_coder. |
| (element_and_restriction, current_watermark, deferred_timestamp) = ( |
| deferred_remainder) |
| if deferred_timestamp: |
| assert isinstance(deferred_timestamp, timestamp.Duration) |
| proto_deferred_watermark = proto_utils.from_micros( |
| duration_pb2.Duration, |
| deferred_timestamp.micros) # type: Optional[duration_pb2.Duration] |
| else: |
| proto_deferred_watermark = None |
| return beam_fn_api_pb2.DelayedBundleApplication( |
| requested_time_delay=proto_deferred_watermark, |
| application=self.construct_bundle_application( |
| op.input_info, current_watermark, element_and_restriction)) |
| |
| def bundle_application(self, |
| op, # type: operations.DoOperation |
| primary # type: SplitResultPrimary |
| ): |
| # type: (...) -> beam_fn_api_pb2.BundleApplication |
| assert op.input_info is not None |
| return self.construct_bundle_application( |
| op.input_info, None, primary.primary_value) |
| |
| def construct_bundle_application(self, |
| op_input_info, # type: operations.OpInputInfo |
| output_watermark, # type: Optional[timestamp.Timestamp] |
| element |
| ): |
| # type: (...) -> beam_fn_api_pb2.BundleApplication |
| transform_id, main_input_tag, main_input_coder, outputs = op_input_info |
| if output_watermark: |
| proto_output_watermark = proto_utils.from_micros( |
| timestamp_pb2.Timestamp, output_watermark.micros) |
| output_watermarks = { |
| output: proto_output_watermark |
| for output in outputs |
| } # type: Optional[Dict[str, timestamp_pb2.Timestamp]] |
| else: |
| output_watermarks = None |
| return beam_fn_api_pb2.BundleApplication( |
| transform_id=transform_id, |
| input_id=main_input_tag, |
| output_watermarks=output_watermarks, |
| element=main_input_coder.get_impl().encode_nested(element)) |
| |
| def monitoring_infos(self): |
| # type: () -> List[metrics_pb2.MonitoringInfo] |
| |
| """Returns the list of MonitoringInfos collected processing this bundle.""" |
| # Construct a new dict first to remove duplicates. |
| all_monitoring_infos_dict = {} |
| for transform_id, op in self.ops.items(): |
| tag_to_pcollection_id = self.process_bundle_descriptor.transforms[ |
| transform_id].outputs |
| all_monitoring_infos_dict.update( |
| op.monitoring_infos(transform_id, dict(tag_to_pcollection_id))) |
| |
| return list(all_monitoring_infos_dict.values()) |
| |
| def shutdown(self): |
| # type: () -> None |
| for op in self.ops.values(): |
| op.teardown() |
| |
| |
| @dataclass |
| class ExecutionContext: |
| # Any splits to be processed later. |
| delayed_applications: List[Tuple[operations.DoOperation, |
| common.SplitResultResidual]] = field( |
| default_factory=list) |
| |
| # The exception sampler for the currently executing PTransform. |
| output_sampler: Optional[data_sampler.OutputSampler] = None |
| |
| # The current instruction being executed. |
| instruction_id: Optional[str] = None |
| |
| |
| class BeamTransformFactory(object): |
| """Factory for turning transform_protos into executable operations.""" |
| def __init__(self, |
| runner_capabilities, # type: FrozenSet[str] |
| descriptor, # type: beam_fn_api_pb2.ProcessBundleDescriptor |
| data_channel_factory, # type: data_plane.DataChannelFactory |
| counter_factory, # type: counters.CounterFactory |
| state_sampler, # type: statesampler.StateSampler |
| state_handler, # type: sdk_worker.CachingStateHandler |
| data_sampler, # type: Optional[data_sampler.DataSampler] |
| ): |
| self.runner_capabilities = runner_capabilities |
| self.descriptor = descriptor |
| self.data_channel_factory = data_channel_factory |
| self.counter_factory = counter_factory |
| self.state_sampler = state_sampler |
| self.state_handler = state_handler |
| self.context = pipeline_context.PipelineContext( |
| descriptor, |
| iterable_state_read=lambda token, |
| element_coder_impl: _StateBackedIterable( |
| state_handler, |
| beam_fn_api_pb2.StateKey( |
| runner=beam_fn_api_pb2.StateKey.Runner(key=token)), |
| element_coder_impl)) |
| self.data_sampler = data_sampler |
| |
| _known_urns = { |
| } # type: Dict[str, Tuple[ConstructorFn, Union[Type[message.Message], Type[bytes], None]]] |
| |
| @classmethod |
| def register_urn( |
| cls, |
| urn, # type: str |
| parameter_type # type: Optional[Type[T]] |
| ): |
| # type: (...) -> Callable[[Callable[[BeamTransformFactory, str, beam_runner_api_pb2.PTransform, T, Dict[str, List[operations.Operation]]], operations.Operation]], Callable[[BeamTransformFactory, str, beam_runner_api_pb2.PTransform, T, Dict[str, List[operations.Operation]]], operations.Operation]] |
| def wrapper(func): |
| cls._known_urns[urn] = func, parameter_type |
| return func |
| |
| return wrapper |
| |
| def create_operation(self, |
| transform_id, # type: str |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| # type: (...) -> operations.Operation |
| transform_proto = self.descriptor.transforms[transform_id] |
| if not transform_proto.unique_name: |
| _LOGGER.debug("No unique name set for transform %s" % transform_id) |
| transform_proto.unique_name = transform_id |
| creator, parameter_type = self._known_urns[transform_proto.spec.urn] |
| payload = proto_utils.parse_Bytes( |
| transform_proto.spec.payload, parameter_type) |
| return creator(self, transform_id, transform_proto, payload, consumers) |
| |
| def extract_timers_info(self): |
| # type: () -> Dict[Tuple[str, str], TimerInfo] |
| timers_info = {} |
| for transform_id, transform_proto in self.descriptor.transforms.items(): |
| if transform_proto.spec.urn == common_urns.primitives.PAR_DO.urn: |
| pardo_payload = proto_utils.parse_Bytes( |
| transform_proto.spec.payload, beam_runner_api_pb2.ParDoPayload) |
| for (timer_family_id, |
| timer_family_spec) in pardo_payload.timer_family_specs.items(): |
| timer_coder_impl = self.get_coder( |
| timer_family_spec.timer_family_coder_id).get_impl() |
| # The output_stream should be updated when processing a bundle. |
| timers_info[(transform_id, timer_family_id)] = TimerInfo( |
| timer_coder_impl=timer_coder_impl) |
| return timers_info |
| |
| def get_coder(self, coder_id): |
| # type: (str) -> coders.Coder |
| if coder_id not in self.descriptor.coders: |
| raise KeyError("No such coder: %s" % coder_id) |
| coder_proto = self.descriptor.coders[coder_id] |
| if coder_proto.spec.urn: |
| return self.context.coders.get_by_id(coder_id) |
| else: |
| # No URN, assume cloud object encoding json bytes. |
| return operation_specs.get_coder_from_spec( |
| json.loads(coder_proto.spec.payload.decode('utf-8'))) |
| |
| def get_windowed_coder(self, pcoll_id): |
| # type: (str) -> WindowedValueCoder |
| coder = self.get_coder(self.descriptor.pcollections[pcoll_id].coder_id) |
| # TODO(robertwb): Remove this condition once all runners are consistent. |
| if not isinstance(coder, WindowedValueCoder): |
| windowing_strategy = self.descriptor.windowing_strategies[ |
| self.descriptor.pcollections[pcoll_id].windowing_strategy_id] |
| return WindowedValueCoder( |
| coder, self.get_coder(windowing_strategy.window_coder_id)) |
| else: |
| return coder |
| |
| def get_output_coders(self, transform_proto): |
| # type: (beam_runner_api_pb2.PTransform) -> Dict[str, coders.Coder] |
| return { |
| tag: self.get_windowed_coder(pcoll_id) |
| for tag, |
| pcoll_id in transform_proto.outputs.items() |
| } |
| |
| def get_only_output_coder(self, transform_proto): |
| # type: (beam_runner_api_pb2.PTransform) -> coders.Coder |
| return only_element(self.get_output_coders(transform_proto).values()) |
| |
| def get_input_coders(self, transform_proto): |
| # type: (beam_runner_api_pb2.PTransform) -> Dict[str, coders.WindowedValueCoder] |
| return { |
| tag: self.get_windowed_coder(pcoll_id) |
| for tag, |
| pcoll_id in transform_proto.inputs.items() |
| } |
| |
| def get_only_input_coder(self, transform_proto): |
| # type: (beam_runner_api_pb2.PTransform) -> coders.Coder |
| return only_element(list(self.get_input_coders(transform_proto).values())) |
| |
| def get_input_windowing(self, transform_proto): |
| # type: (beam_runner_api_pb2.PTransform) -> Windowing |
| pcoll_id = only_element(transform_proto.inputs.values()) |
| windowing_strategy_id = self.descriptor.pcollections[ |
| pcoll_id].windowing_strategy_id |
| return self.context.windowing_strategies.get_by_id(windowing_strategy_id) |
| |
| # TODO(robertwb): Update all operations to take these in the constructor. |
| @staticmethod |
| def augment_oldstyle_op( |
| op, # type: OperationT |
| step_name, # type: str |
| consumers, # type: Mapping[str, Iterable[operations.Operation]] |
| tag_list=None # type: Optional[List[str]] |
| ): |
| # type: (...) -> OperationT |
| op.step_name = step_name |
| for tag, op_consumers in consumers.items(): |
| for consumer in op_consumers: |
| op.add_receiver(consumer, tag_list.index(tag) if tag_list else 0) |
| return op |
| |
| |
| @BeamTransformFactory.register_urn( |
| DATA_INPUT_URN, beam_fn_api_pb2.RemoteGrpcPort) |
| def create_source_runner( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| grpc_port, # type: beam_fn_api_pb2.RemoteGrpcPort |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| # type: (...) -> DataInputOperation |
| |
| output_coder = factory.get_coder(grpc_port.coder_id) |
| return DataInputOperation( |
| common.NameContext(transform_proto.unique_name, transform_id), |
| transform_proto.unique_name, |
| consumers, |
| factory.counter_factory, |
| factory.state_sampler, |
| output_coder, |
| transform_id=transform_id, |
| data_channel=factory.data_channel_factory.create_data_channel(grpc_port)) |
| |
| |
| @BeamTransformFactory.register_urn( |
| DATA_OUTPUT_URN, beam_fn_api_pb2.RemoteGrpcPort) |
| def create_sink_runner( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| grpc_port, # type: beam_fn_api_pb2.RemoteGrpcPort |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| # type: (...) -> DataOutputOperation |
| output_coder = factory.get_coder(grpc_port.coder_id) |
| return DataOutputOperation( |
| common.NameContext(transform_proto.unique_name, transform_id), |
| transform_proto.unique_name, |
| consumers, |
| factory.counter_factory, |
| factory.state_sampler, |
| output_coder, |
| transform_id=transform_id, |
| data_channel=factory.data_channel_factory.create_data_channel(grpc_port)) |
| |
| |
| @BeamTransformFactory.register_urn(OLD_DATAFLOW_RUNNER_HARNESS_READ_URN, None) |
| def create_source_java( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| parameter, |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| # type: (...) -> operations.ReadOperation |
| # The Dataflow runner harness strips the base64 encoding. |
| source = pickler.loads(base64.b64encode(parameter)) |
| spec = operation_specs.WorkerRead( |
| iobase.SourceBundle(1.0, source, None, None), |
| [factory.get_only_output_coder(transform_proto)]) |
| return factory.augment_oldstyle_op( |
| operations.ReadOperation( |
| common.NameContext(transform_proto.unique_name, transform_id), |
| spec, |
| factory.counter_factory, |
| factory.state_sampler), |
| transform_proto.unique_name, |
| consumers) |
| |
| |
| @BeamTransformFactory.register_urn( |
| common_urns.deprecated_primitives.READ.urn, beam_runner_api_pb2.ReadPayload) |
| def create_deprecated_read( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| parameter, # type: beam_runner_api_pb2.ReadPayload |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| # type: (...) -> operations.ReadOperation |
| source = iobase.BoundedSource.from_runner_api( |
| parameter.source, factory.context) |
| spec = operation_specs.WorkerRead( |
| iobase.SourceBundle(1.0, source, None, None), |
| [WindowedValueCoder(source.default_output_coder())]) |
| return factory.augment_oldstyle_op( |
| operations.ReadOperation( |
| common.NameContext(transform_proto.unique_name, transform_id), |
| spec, |
| factory.counter_factory, |
| factory.state_sampler), |
| transform_proto.unique_name, |
| consumers) |
| |
| |
| @BeamTransformFactory.register_urn( |
| python_urns.IMPULSE_READ_TRANSFORM, beam_runner_api_pb2.ReadPayload) |
| def create_read_from_impulse_python( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| parameter, # type: beam_runner_api_pb2.ReadPayload |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| # type: (...) -> operations.ImpulseReadOperation |
| return operations.ImpulseReadOperation( |
| common.NameContext(transform_proto.unique_name, transform_id), |
| factory.counter_factory, |
| factory.state_sampler, |
| consumers, |
| iobase.BoundedSource.from_runner_api(parameter.source, factory.context), |
| factory.get_only_output_coder(transform_proto)) |
| |
| |
| @BeamTransformFactory.register_urn(OLD_DATAFLOW_RUNNER_HARNESS_PARDO_URN, None) |
| def create_dofn_javasdk( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| serialized_fn, |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| return _create_pardo_operation( |
| factory, transform_id, transform_proto, consumers, serialized_fn) |
| |
| |
| @BeamTransformFactory.register_urn( |
| common_urns.sdf_components.PAIR_WITH_RESTRICTION.urn, |
| beam_runner_api_pb2.ParDoPayload) |
| def create_pair_with_restriction(*args): |
| class PairWithRestriction(beam.DoFn): |
| def __init__(self, fn, restriction_provider, watermark_estimator_provider): |
| self.restriction_provider = restriction_provider |
| self.watermark_estimator_provider = watermark_estimator_provider |
| |
| def process(self, element, *args, **kwargs): |
| # TODO(SDF): Do we want to allow mutation of the element? |
| # (E.g. it could be nice to shift bulky description to the portion |
| # that can be distributed.) |
| initial_restriction = self.restriction_provider.initial_restriction( |
| element) |
| initial_estimator_state = ( |
| self.watermark_estimator_provider.initial_estimator_state( |
| element, initial_restriction)) |
| yield (element, (initial_restriction, initial_estimator_state)) |
| |
| return _create_sdf_operation(PairWithRestriction, *args) |
| |
| |
| @BeamTransformFactory.register_urn( |
| common_urns.sdf_components.SPLIT_AND_SIZE_RESTRICTIONS.urn, |
| beam_runner_api_pb2.ParDoPayload) |
| def create_split_and_size_restrictions(*args): |
| class SplitAndSizeRestrictions(beam.DoFn): |
| def __init__(self, fn, restriction_provider, watermark_estimator_provider): |
| self.restriction_provider = restriction_provider |
| self.watermark_estimator_provider = watermark_estimator_provider |
| |
| def process(self, element_restriction, *args, **kwargs): |
| element, (restriction, _) = element_restriction |
| for part, size in self.restriction_provider.split_and_size( |
| element, restriction): |
| if size < 0: |
| raise ValueError('Expected size >= 0 but received %s.' % size) |
| estimator_state = ( |
| self.watermark_estimator_provider.initial_estimator_state( |
| element, part)) |
| yield ((element, (part, estimator_state)), size) |
| |
| return _create_sdf_operation(SplitAndSizeRestrictions, *args) |
| |
| |
| @BeamTransformFactory.register_urn( |
| common_urns.sdf_components.TRUNCATE_SIZED_RESTRICTION.urn, |
| beam_runner_api_pb2.ParDoPayload) |
| def create_truncate_sized_restriction(*args): |
| class TruncateAndSizeRestriction(beam.DoFn): |
| def __init__(self, fn, restriction_provider, watermark_estimator_provider): |
| self.restriction_provider = restriction_provider |
| |
| def process(self, element_restriction, *args, **kwargs): |
| ((element, (restriction, estimator_state)), _) = element_restriction |
| truncated_restriction = self.restriction_provider.truncate( |
| element, restriction) |
| if truncated_restriction: |
| truncated_restriction_size = ( |
| self.restriction_provider.restriction_size( |
| element, truncated_restriction)) |
| if truncated_restriction_size < 0: |
| raise ValueError( |
| 'Expected size >= 0 but received %s.' % |
| truncated_restriction_size) |
| yield ((element, (truncated_restriction, estimator_state)), |
| truncated_restriction_size) |
| |
| return _create_sdf_operation( |
| TruncateAndSizeRestriction, |
| *args, |
| operation_cls=operations.SdfTruncateSizedRestrictions) |
| |
| |
| @BeamTransformFactory.register_urn( |
| common_urns.sdf_components.PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS.urn, |
| beam_runner_api_pb2.ParDoPayload) |
| def create_process_sized_elements_and_restrictions( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| parameter, # type: beam_runner_api_pb2.ParDoPayload |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| return _create_pardo_operation( |
| factory, |
| transform_id, |
| transform_proto, |
| consumers, |
| core.DoFnInfo.from_runner_api(parameter.do_fn, |
| factory.context).serialized_dofn_data(), |
| parameter, |
| operation_cls=operations.SdfProcessSizedElements) |
| |
| |
| def _create_sdf_operation( |
| proxy_dofn, |
| factory, |
| transform_id, |
| transform_proto, |
| parameter, |
| consumers, |
| operation_cls=operations.DoOperation): |
| |
| dofn_data = pickler.loads(parameter.do_fn.payload) |
| dofn = dofn_data[0] |
| restriction_provider = common.DoFnSignature(dofn).get_restriction_provider() |
| watermark_estimator_provider = ( |
| common.DoFnSignature(dofn).get_watermark_estimator_provider()) |
| serialized_fn = pickler.dumps( |
| (proxy_dofn(dofn, restriction_provider, watermark_estimator_provider), ) + |
| dofn_data[1:]) |
| return _create_pardo_operation( |
| factory, |
| transform_id, |
| transform_proto, |
| consumers, |
| serialized_fn, |
| parameter, |
| operation_cls=operation_cls) |
| |
| |
| @BeamTransformFactory.register_urn( |
| common_urns.primitives.PAR_DO.urn, beam_runner_api_pb2.ParDoPayload) |
| def create_par_do( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| parameter, # type: beam_runner_api_pb2.ParDoPayload |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| # type: (...) -> operations.DoOperation |
| return _create_pardo_operation( |
| factory, |
| transform_id, |
| transform_proto, |
| consumers, |
| core.DoFnInfo.from_runner_api(parameter.do_fn, |
| factory.context).serialized_dofn_data(), |
| parameter) |
| |
| |
| def _create_pardo_operation( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| consumers, |
| serialized_fn, |
| pardo_proto=None, # type: Optional[beam_runner_api_pb2.ParDoPayload] |
| operation_cls=operations.DoOperation |
| ): |
| |
| if pardo_proto and pardo_proto.side_inputs: |
| input_tags_to_coders = factory.get_input_coders(transform_proto) |
| tagged_side_inputs = [ |
| (tag, beam.pvalue.SideInputData.from_runner_api(si, factory.context)) |
| for tag, |
| si in pardo_proto.side_inputs.items() |
| ] |
| tagged_side_inputs.sort( |
| key=lambda tag_si: sideinputs.get_sideinput_index(tag_si[0])) |
| side_input_maps = [ |
| StateBackedSideInputMap( |
| factory.state_handler, |
| transform_id, |
| tag, |
| si, |
| input_tags_to_coders[tag], |
| use_bulk_read=( |
| common_urns.runner_protocols.MULTIMAP_KEYS_VALUES_SIDE_INPUT.urn |
| in factory.runner_capabilities)) |
| for (tag, si) in tagged_side_inputs |
| ] |
| else: |
| side_input_maps = [] |
| |
| output_tags = list(transform_proto.outputs.keys()) |
| |
| dofn_data = pickler.loads(serialized_fn) |
| if not dofn_data[-1]: |
| # Windowing not set. |
| if pardo_proto: |
| other_input_tags = set.union( |
| set(pardo_proto.side_inputs), |
| set(pardo_proto.timer_family_specs)) # type: Container[str] |
| else: |
| other_input_tags = () |
| pcoll_id, = [pcoll for tag, pcoll in transform_proto.inputs.items() |
| if tag not in other_input_tags] |
| windowing = factory.context.windowing_strategies.get_by_id( |
| factory.descriptor.pcollections[pcoll_id].windowing_strategy_id) |
| serialized_fn = pickler.dumps(dofn_data[:-1] + (windowing, )) |
| |
| if pardo_proto and (pardo_proto.timer_family_specs or pardo_proto.state_specs |
| or pardo_proto.restriction_coder_id): |
| found_input_coder = None |
| for tag, pcoll_id in transform_proto.inputs.items(): |
| if tag in pardo_proto.side_inputs: |
| pass |
| else: |
| # Must be the main input |
| assert found_input_coder is None |
| main_input_tag = tag |
| found_input_coder = factory.get_windowed_coder(pcoll_id) |
| assert found_input_coder is not None |
| main_input_coder = found_input_coder |
| |
| if pardo_proto.timer_family_specs or pardo_proto.state_specs: |
| user_state_context = FnApiUserStateContext( |
| factory.state_handler, |
| transform_id, |
| main_input_coder.key_coder(), |
| main_input_coder.window_coder |
| ) # type: Optional[FnApiUserStateContext] |
| else: |
| user_state_context = None |
| else: |
| user_state_context = None |
| |
| output_coders = factory.get_output_coders(transform_proto) |
| spec = operation_specs.WorkerDoFn( |
| serialized_fn=serialized_fn, |
| output_tags=output_tags, |
| input=None, |
| side_inputs=None, # Fn API uses proto definitions and the Fn State API |
| output_coders=[output_coders[tag] for tag in output_tags]) |
| |
| result = factory.augment_oldstyle_op( |
| operation_cls( |
| common.NameContext(transform_proto.unique_name, transform_id), |
| spec, |
| factory.counter_factory, |
| factory.state_sampler, |
| side_input_maps, |
| user_state_context), |
| transform_proto.unique_name, |
| consumers, |
| output_tags) |
| if pardo_proto and pardo_proto.restriction_coder_id: |
| result.input_info = operations.OpInputInfo( |
| transform_id, |
| main_input_tag, |
| main_input_coder, |
| transform_proto.outputs.keys()) |
| return result |
| |
| |
| def _create_simple_pardo_operation(factory, # type: BeamTransformFactory |
| transform_id, |
| transform_proto, |
| consumers, |
| dofn, # type: beam.DoFn |
| ): |
| serialized_fn = pickler.dumps((dofn, (), {}, [], None)) |
| return _create_pardo_operation( |
| factory, transform_id, transform_proto, consumers, serialized_fn) |
| |
| |
| @BeamTransformFactory.register_urn( |
| common_urns.primitives.ASSIGN_WINDOWS.urn, |
| beam_runner_api_pb2.WindowingStrategy) |
| def create_assign_windows( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| parameter, # type: beam_runner_api_pb2.WindowingStrategy |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| class WindowIntoDoFn(beam.DoFn): |
| def __init__(self, windowing): |
| self.windowing = windowing |
| |
| def process( |
| self, |
| element, |
| timestamp=beam.DoFn.TimestampParam, |
| window=beam.DoFn.WindowParam): |
| new_windows = self.windowing.windowfn.assign( |
| WindowFn.AssignContext(timestamp, element=element, window=window)) |
| yield WindowedValue(element, timestamp, new_windows) |
| |
| from apache_beam.transforms.core import Windowing |
| from apache_beam.transforms.window import WindowFn |
| windowing = Windowing.from_runner_api(parameter, factory.context) |
| return _create_simple_pardo_operation( |
| factory, |
| transform_id, |
| transform_proto, |
| consumers, |
| WindowIntoDoFn(windowing)) |
| |
| |
| @BeamTransformFactory.register_urn(IDENTITY_DOFN_URN, None) |
| def create_identity_dofn( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| parameter, |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| # type: (...) -> operations.FlattenOperation |
| return factory.augment_oldstyle_op( |
| operations.FlattenOperation( |
| common.NameContext(transform_proto.unique_name, transform_id), |
| operation_specs.WorkerFlatten( |
| None, [factory.get_only_output_coder(transform_proto)]), |
| factory.counter_factory, |
| factory.state_sampler), |
| transform_proto.unique_name, |
| consumers) |
| |
| |
| @BeamTransformFactory.register_urn( |
| common_urns.combine_components.COMBINE_PER_KEY_PRECOMBINE.urn, |
| beam_runner_api_pb2.CombinePayload) |
| def create_combine_per_key_precombine( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| payload, # type: beam_runner_api_pb2.CombinePayload |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| # type: (...) -> operations.PGBKCVOperation |
| serialized_combine_fn = pickler.dumps(( |
| beam.CombineFn.from_runner_api(payload.combine_fn, |
| factory.context), [], {})) |
| return factory.augment_oldstyle_op( |
| operations.PGBKCVOperation( |
| common.NameContext(transform_proto.unique_name, transform_id), |
| operation_specs.WorkerPartialGroupByKey( |
| serialized_combine_fn, |
| None, [factory.get_only_output_coder(transform_proto)]), |
| factory.counter_factory, |
| factory.state_sampler, |
| factory.get_input_windowing(transform_proto)), |
| transform_proto.unique_name, |
| consumers) |
| |
| |
| @BeamTransformFactory.register_urn( |
| common_urns.combine_components.COMBINE_PER_KEY_MERGE_ACCUMULATORS.urn, |
| beam_runner_api_pb2.CombinePayload) |
| def create_combbine_per_key_merge_accumulators( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| payload, # type: beam_runner_api_pb2.CombinePayload |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| return _create_combine_phase_operation( |
| factory, transform_id, transform_proto, payload, consumers, 'merge') |
| |
| |
| @BeamTransformFactory.register_urn( |
| common_urns.combine_components.COMBINE_PER_KEY_EXTRACT_OUTPUTS.urn, |
| beam_runner_api_pb2.CombinePayload) |
| def create_combine_per_key_extract_outputs( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| payload, # type: beam_runner_api_pb2.CombinePayload |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| return _create_combine_phase_operation( |
| factory, transform_id, transform_proto, payload, consumers, 'extract') |
| |
| |
| @BeamTransformFactory.register_urn( |
| common_urns.combine_components.COMBINE_PER_KEY_CONVERT_TO_ACCUMULATORS.urn, |
| beam_runner_api_pb2.CombinePayload) |
| def create_combine_per_key_convert_to_accumulators( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| payload, # type: beam_runner_api_pb2.CombinePayload |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| return _create_combine_phase_operation( |
| factory, transform_id, transform_proto, payload, consumers, 'convert') |
| |
| |
| @BeamTransformFactory.register_urn( |
| common_urns.combine_components.COMBINE_GROUPED_VALUES.urn, |
| beam_runner_api_pb2.CombinePayload) |
| def create_combine_grouped_values( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| payload, # type: beam_runner_api_pb2.CombinePayload |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| return _create_combine_phase_operation( |
| factory, transform_id, transform_proto, payload, consumers, 'all') |
| |
| |
| def _create_combine_phase_operation( |
| factory, transform_id, transform_proto, payload, consumers, phase): |
| # type: (...) -> operations.CombineOperation |
| serialized_combine_fn = pickler.dumps(( |
| beam.CombineFn.from_runner_api(payload.combine_fn, |
| factory.context), [], {})) |
| return factory.augment_oldstyle_op( |
| operations.CombineOperation( |
| common.NameContext(transform_proto.unique_name, transform_id), |
| operation_specs.WorkerCombineFn( |
| serialized_combine_fn, |
| phase, |
| None, [factory.get_only_output_coder(transform_proto)]), |
| factory.counter_factory, |
| factory.state_sampler), |
| transform_proto.unique_name, |
| consumers) |
| |
| |
| @BeamTransformFactory.register_urn(common_urns.primitives.FLATTEN.urn, None) |
| def create_flatten( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| payload, |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| # type: (...) -> operations.FlattenOperation |
| return factory.augment_oldstyle_op( |
| operations.FlattenOperation( |
| common.NameContext(transform_proto.unique_name, transform_id), |
| operation_specs.WorkerFlatten( |
| None, [factory.get_only_output_coder(transform_proto)]), |
| factory.counter_factory, |
| factory.state_sampler), |
| transform_proto.unique_name, |
| consumers) |
| |
| |
| @BeamTransformFactory.register_urn( |
| common_urns.primitives.MAP_WINDOWS.urn, beam_runner_api_pb2.FunctionSpec) |
| def create_map_windows( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| mapping_fn_spec, # type: beam_runner_api_pb2.FunctionSpec |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| assert mapping_fn_spec.urn == python_urns.PICKLED_WINDOW_MAPPING_FN |
| window_mapping_fn = pickler.loads(mapping_fn_spec.payload) |
| |
| class MapWindows(beam.DoFn): |
| def process(self, element): |
| key, window = element |
| return [(key, window_mapping_fn(window))] |
| |
| return _create_simple_pardo_operation( |
| factory, transform_id, transform_proto, consumers, MapWindows()) |
| |
| |
| @BeamTransformFactory.register_urn( |
| common_urns.primitives.MERGE_WINDOWS.urn, beam_runner_api_pb2.FunctionSpec) |
| def create_merge_windows( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| mapping_fn_spec, # type: beam_runner_api_pb2.FunctionSpec |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| assert mapping_fn_spec.urn == python_urns.PICKLED_WINDOWFN |
| window_fn = pickler.loads(mapping_fn_spec.payload) |
| |
| class MergeWindows(beam.DoFn): |
| def process(self, element): |
| nonce, windows = element |
| |
| original_windows = set(windows) # type: Set[window.BoundedWindow] |
| merged_windows = collections.defaultdict( |
| set |
| ) # type: MutableMapping[window.BoundedWindow, Set[window.BoundedWindow]] # noqa: F821 |
| |
| class RecordingMergeContext(window.WindowFn.MergeContext): |
| def merge( |
| self, |
| to_be_merged, # type: Iterable[window.BoundedWindow] |
| merge_result, # type: window.BoundedWindow |
| ): |
| originals = merged_windows[merge_result] |
| for window in to_be_merged: |
| if window in original_windows: |
| originals.add(window) |
| original_windows.remove(window) |
| else: |
| originals.update(merged_windows.pop(window)) |
| |
| window_fn.merge(RecordingMergeContext(windows)) |
| yield nonce, (original_windows, merged_windows.items()) |
| |
| return _create_simple_pardo_operation( |
| factory, transform_id, transform_proto, consumers, MergeWindows()) |
| |
| |
| @BeamTransformFactory.register_urn(common_urns.primitives.TO_STRING.urn, None) |
| def create_to_string_fn( |
| factory, # type: BeamTransformFactory |
| transform_id, # type: str |
| transform_proto, # type: beam_runner_api_pb2.PTransform |
| mapping_fn_spec, # type: beam_runner_api_pb2.FunctionSpec |
| consumers # type: Dict[str, List[operations.Operation]] |
| ): |
| class ToString(beam.DoFn): |
| def process(self, element): |
| key, value = element |
| return [(key, str(value))] |
| |
| return _create_simple_pardo_operation( |
| factory, transform_id, transform_proto, consumers, ToString()) |
