You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by ro...@apache.org on 2017/10/22 07:40:37 UTC
[2/2] beam git commit: Implement FnApi side inputs in Python.
Implement FnApi side inputs in Python.
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/38556b78
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/38556b78
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/38556b78
Branch: refs/heads/master
Commit: 38556b78bfa874620ccb61a5179a44c7a51dbf55
Parents: a71f000
Author: Robert Bradshaw <ro...@gmail.com>
Authored: Thu Oct 19 12:11:44 2017 -0700
Committer: Robert Bradshaw <ro...@gmail.com>
Committed: Sun Oct 22 00:37:19 2017 -0700
----------------------------------------------------------------------
sdks/python/apache_beam/pipeline.py | 30 ++--
sdks/python/apache_beam/pvalue.py | 103 +++++++++++++-
.../runners/direct/transform_evaluator.py | 9 +-
.../runners/portability/fn_api_runner.py | 142 +++++++++++++++----
.../runners/portability/fn_api_runner_test.py | 41 +++++-
.../runners/worker/bundle_processor.py | 116 +++++++++------
.../apache_beam/runners/worker/data_plane.py | 18 ++-
.../apache_beam/runners/worker/operations.pxd | 1 +
.../apache_beam/runners/worker/operations.py | 17 ++-
.../apache_beam/runners/worker/sdk_worker.py | 111 ++++++++++++++-
sdks/python/apache_beam/transforms/core.py | 23 ++-
sdks/python/apache_beam/utils/urns.py | 5 +
12 files changed, 511 insertions(+), 105 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/38556b78/sdks/python/apache_beam/pipeline.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/pipeline.py b/sdks/python/apache_beam/pipeline.py
index c670978..62626a3 100644
--- a/sdks/python/apache_beam/pipeline.py
+++ b/sdks/python/apache_beam/pipeline.py
@@ -506,9 +506,6 @@ class Pipeline(object):
self.visit_transform(transform_node)
def visit_transform(self, transform_node):
- if transform_node.side_inputs:
- # No side inputs (yet).
- Visitor.ok = False
try:
# Transforms must be picklable.
pickler.loads(pickler.dumps(transform_node.transform,
@@ -730,8 +727,12 @@ class AppliedPTransform(object):
def named_inputs(self):
# TODO(BEAM-1833): Push names up into the sdk construction.
- return {str(ix): input for ix, input in enumerate(self.inputs)
- if isinstance(input, pvalue.PCollection)}
+ main_inputs = {str(ix): input
+ for ix, input in enumerate(self.inputs)
+ if isinstance(input, pvalue.PCollection)}
+ side_inputs = {'side%s' % ix: si.pvalue
+ for ix, si in enumerate(self.side_inputs)}
+ return dict(main_inputs, **side_inputs)
def named_outputs(self):
return {str(tag): output for tag, output in self.outputs.items()
@@ -750,7 +751,6 @@ class AppliedPTransform(object):
spec=transform_to_runner_api(self.transform, context),
subtransforms=[context.transforms.get_id(part, label=part.full_label)
for part in self.parts],
- # TODO(BEAM-115): Side inputs.
inputs={tag: context.pcollections.get_id(pc)
for tag, pc in self.named_inputs().items()},
outputs={str(tag): context.pcollections.get_id(out)
@@ -760,12 +760,26 @@ class AppliedPTransform(object):
@staticmethod
def from_runner_api(proto, context):
+ def is_side_input(tag):
+ # As per named_inputs() above.
+ return tag.startswith('side')
+ main_inputs = [context.pcollections.get_by_id(id)
+ for tag, id in proto.inputs.items()
+ if not is_side_input(tag)]
+ # Ordering is important here.
+ indexed_side_inputs = [(int(tag[4:]), context.pcollections.get_by_id(id))
+ for tag, id in proto.inputs.items()
+ if is_side_input(tag)]
+ side_inputs = [si for _, si in sorted(indexed_side_inputs)]
result = AppliedPTransform(
parent=None,
transform=ptransform.PTransform.from_runner_api(proto.spec, context),
full_label=proto.unique_name,
- inputs=[
- context.pcollections.get_by_id(id) for id in proto.inputs.values()])
+ inputs=main_inputs)
+ if result.transform and result.transform.side_inputs:
+ for si, pcoll in zip(result.transform.side_inputs, side_inputs):
+ si.pvalue = pcoll
+ result.side_inputs = tuple(result.transform.side_inputs)
result.parts = [
context.transforms.get_by_id(id) for id in proto.subtransforms]
result.outputs = {
http://git-wip-us.apache.org/repos/asf/beam/blob/38556b78/sdks/python/apache_beam/pvalue.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/pvalue.py b/sdks/python/apache_beam/pvalue.py
index d2d3653..31922f3 100644
--- a/sdks/python/apache_beam/pvalue.py
+++ b/sdks/python/apache_beam/pvalue.py
@@ -28,7 +28,11 @@ from __future__ import absolute_import
import itertools
+from apache_beam import coders
from apache_beam import typehints
+from apache_beam.internal import pickler
+from apache_beam.portability.api import beam_runner_api_pb2
+from apache_beam.utils import urns
__all__ = [
'PCollection',
@@ -127,8 +131,6 @@ class PCollection(PValue):
return _InvalidUnpickledPCollection, ()
def to_runner_api(self, context):
- from apache_beam.portability.api import beam_runner_api_pb2
- from apache_beam.internal import pickler
return beam_runner_api_pb2.PCollection(
unique_name='%d%s.%s' % (
len(self.producer.full_label), self.producer.full_label, self.tag),
@@ -139,7 +141,6 @@ class PCollection(PValue):
@staticmethod
def from_runner_api(proto, context):
- from apache_beam.internal import pickler
# Producer and tag will be filled in later, the key point is that the
# same object is returned for the same pcollection id.
return PCollection(None, element_type=pickler.loads(proto.coder_id))
@@ -288,6 +289,81 @@ class AsSideInput(object):
def element_type(self):
return typehints.Any
+ # TODO(robertwb): Get rid of _from_runtime_iterable and _view_options
+ # in favor of _side_input_data().
+ def _side_input_data(self):
+ view_options = self._view_options()
+ from_runtime_iterable = type(self)._from_runtime_iterable
+ return SideInputData(
+ urns.ITERABLE_ACCESS,
+ self._window_mapping_fn,
+ lambda iterable: from_runtime_iterable(iterable, view_options),
+ self._input_element_coder())
+
+ def _input_element_coder(self):
+ return coders.WindowedValueCoder(
+ coders.registry.get_coder(self.pvalue.element_type),
+ window_coder=self.pvalue.windowing.windowfn.get_window_coder())
+
+ def to_runner_api(self, context):
+ return self._side_input_data().to_runner_api(context)
+
+ @staticmethod
+ def from_runner_api(proto, context):
+ return _UnpickledSideInput(
+ SideInputData.from_runner_api(proto, context))
+
+
+class _UnpickledSideInput(AsSideInput):
+ def __init__(self, side_input_data):
+ self._data = side_input_data
+ self._window_mapping_fn = side_input_data.window_mapping_fn
+
+ @staticmethod
+ def _from_runtime_iterable(it, options):
+ return options['data'].view_fn(it)
+
+ def _view_options(self):
+ return {
+ 'data': self._data,
+ # For non-fn-api runners.
+ 'window_mapping_fn': self._data.window_mapping_fn,
+ }
+
+ def _side_input_data(self):
+ return self._data
+
+
+class SideInputData(object):
+ """All of the data about a side input except for the bound PCollection."""
+ def __init__(self, access_pattern, window_mapping_fn, view_fn, coder):
+ self.access_pattern = access_pattern
+ self.window_mapping_fn = window_mapping_fn
+ self.view_fn = view_fn
+ self.coder = coder
+
+ def to_runner_api(self, unused_context):
+ return beam_runner_api_pb2.SideInput(
+ access_pattern=beam_runner_api_pb2.FunctionSpec(
+ urn=self.access_pattern),
+ view_fn=beam_runner_api_pb2.SdkFunctionSpec(
+ spec=beam_runner_api_pb2.FunctionSpec(
+ urn=urns.PICKLED_PYTHON_VIEWFN,
+ payload=pickler.dumps((self.view_fn, self.coder)))),
+ window_mapping_fn=beam_runner_api_pb2.SdkFunctionSpec(
+ spec=beam_runner_api_pb2.FunctionSpec(
+ urn=urns.PICKLED_WINDOW_MAPPING_FN,
+ payload=pickler.dumps(self.window_mapping_fn))))
+
+ @staticmethod
+ def from_runner_api(proto, unused_context):
+ assert proto.view_fn.spec.urn == urns.PICKLED_PYTHON_VIEWFN
+ assert proto.window_mapping_fn.spec.urn == urns.PICKLED_WINDOW_MAPPING_FN
+ return SideInputData(
+ proto.access_pattern.urn,
+ pickler.loads(proto.window_mapping_fn.spec.payload),
+ *pickler.loads(proto.view_fn.spec.payload))
+
class AsSingleton(AsSideInput):
"""Marker specifying that an entire PCollection is to be used as a side input.
@@ -358,6 +434,13 @@ class AsIter(AsSideInput):
def _from_runtime_iterable(it, options):
return it
+ def _side_input_data(self):
+ return SideInputData(
+ urns.ITERABLE_ACCESS,
+ self._window_mapping_fn,
+ lambda iterable: iterable,
+ self._input_element_coder())
+
@property
def element_type(self):
return typehints.Iterable[self.pvalue.element_type]
@@ -382,6 +465,13 @@ class AsList(AsSideInput):
def _from_runtime_iterable(it, options):
return list(it)
+ def _side_input_data(self):
+ return SideInputData(
+ urns.ITERABLE_ACCESS,
+ self._window_mapping_fn,
+ list,
+ self._input_element_coder())
+
class AsDict(AsSideInput):
"""Marker specifying a PCollection to be used as an indexable side input.
@@ -403,6 +493,13 @@ class AsDict(AsSideInput):
def _from_runtime_iterable(it, options):
return dict(it)
+ def _side_input_data(self):
+ return SideInputData(
+ urns.ITERABLE_ACCESS,
+ self._window_mapping_fn,
+ dict,
+ self._input_element_coder())
+
class EmptySideInput(object):
"""Value indicating when a singleton side input was empty.
http://git-wip-us.apache.org/repos/asf/beam/blob/38556b78/sdks/python/apache_beam/runners/direct/transform_evaluator.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/direct/transform_evaluator.py b/sdks/python/apache_beam/runners/direct/transform_evaluator.py
index 16a2991..2f3ac4f 100644
--- a/sdks/python/apache_beam/runners/direct/transform_evaluator.py
+++ b/sdks/python/apache_beam/runners/direct/transform_evaluator.py
@@ -584,10 +584,13 @@ class _GroupByKeyOnlyEvaluator(_TransformEvaluator):
assert len(self._outputs) == 1
self.output_pcollection = list(self._outputs)[0]
- # The input type of a GroupByKey will be KV[Any, Any] or more specific.
+ # The output type of a GroupByKey will be KV[Any, Any] or more specific.
+ # TODO(BEAM-2717): Infer coders earlier.
kv_type_hint = (
- self._applied_ptransform.transform.get_type_hints().input_types[0])
- self.key_coder = coders.registry.get_coder(kv_type_hint[0].tuple_types[0])
+ self._applied_ptransform.outputs[None].element_type
+ or
+ self._applied_ptransform.transform.get_type_hints().input_types[0][0])
+ self.key_coder = coders.registry.get_coder(kv_type_hint.tuple_types[0])
def process_timer(self, timer_firing):
# We do not need to emit a KeyedWorkItem to process_element().
http://git-wip-us.apache.org/repos/asf/beam/blob/38556b78/sdks/python/apache_beam/runners/portability/fn_api_runner.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner.py b/sdks/python/apache_beam/runners/portability/fn_api_runner.py
index 463f78f..838ce1e 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner.py
@@ -65,7 +65,8 @@ def streaming_rpc_handler(cls, method_name):
self._push_queue = queue.Queue()
self._pull_queue = queue.Queue()
setattr(self, method_name, self.run)
- self._read_thread = threading.Thread(target=self._read)
+ self._read_thread = threading.Thread(
+ name='streaming_rpc_handler_read', target=self._read)
self._started = False
def run(self, iterator, context):
@@ -155,6 +156,35 @@ class _GroupingBuffer(object):
return iter([output_stream.get()])
+class _WindowGroupingBuffer(object):
+ """Used to partition windowed side inputs."""
+ def __init__(self, side_input_data):
+ # Here's where we would use a different type of partitioning
+ # (e.g. also by key) for a different access pattern.
+ assert side_input_data.access_pattern == urns.ITERABLE_ACCESS
+ self._windowed_value_coder = side_input_data.coder
+ self._window_coder = side_input_data.coder.window_coder
+ self._value_coder = side_input_data.coder.wrapped_value_coder
+ self._values_by_window = collections.defaultdict(list)
+
+ def append(self, elements_data):
+ input_stream = create_InputStream(elements_data)
+ while input_stream.size() > 0:
+ windowed_value = self._windowed_value_coder.get_impl(
+ ).decode_from_stream(input_stream, True)
+ for window in windowed_value.windows:
+ self._values_by_window[window].append(windowed_value.value)
+
+ def items(self):
+ value_coder_impl = self._value_coder.get_impl()
+ for window, values in self._values_by_window.items():
+ encoded_window = self._window_coder.encode(window)
+ output_stream = create_OutputStream()
+ for value in values:
+ value_coder_impl.encode_to_stream(value, output_stream, True)
+ yield encoded_window, output_stream.get()
+
+
class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
def __init__(self, use_grpc=False, sdk_harness_factory=None):
@@ -206,11 +236,14 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
def __repr__(self):
must_follow = ', '.join(prev.name for prev in self.must_follow)
- return "%s\n %s\n must follow: %s" % (
+ downstream_side_inputs = ', '.join(
+ str(si) for si in self.downstream_side_inputs)
+ return "%s\n %s\n must follow: %s\n downstream_side_inputs: %s" % (
self.name,
'\n'.join(["%s:%s" % (transform.unique_name, transform.spec.urn)
for transform in self.transforms]),
- must_follow)
+ must_follow,
+ downstream_side_inputs)
def can_fuse(self, consumer):
def no_overlap(a, b):
@@ -469,11 +502,12 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
for transform in stage.transforms:
for output in transform.outputs.values():
if output in all_side_inputs:
- downstream_side_inputs = union(downstream_side_inputs, output)
- for consumer in consumers[output]:
- downstream_side_inputs = union(
- downstream_side_inputs,
- compute_downstream_side_inputs(consumer))
+ downstream_side_inputs = union(
+ downstream_side_inputs, frozenset([output]))
+ for consumer in consumers[output]:
+ downstream_side_inputs = union(
+ downstream_side_inputs,
+ compute_downstream_side_inputs(consumer))
downstream_side_inputs_by_stage[stage] = downstream_side_inputs
return downstream_side_inputs_by_stage[stage]
@@ -524,7 +558,7 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
producer = replacement(producer)
consumer = replacement(consumer)
# Update consumer.must_follow set, as it's used in can_fuse.
- consumer.must_follow = set(
+ consumer.must_follow = frozenset(
replacement(s) for s in consumer.must_follow)
if producer.can_fuse(consumer):
fuse(producer, consumer)
@@ -549,8 +583,11 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
spec=beam_runner_api_pb2.FunctionSpec(
urn=bundle_processor.DATA_INPUT_URN,
payload=pcoll_as_param))],
- must_follow={write_pcoll})
+ must_follow=frozenset([write_pcoll]))
fuse(read_pcoll, consumer)
+ else:
+ consumer.must_follow = union(
+ consumer.must_follow, frozenset([write_pcoll]))
# Everything that was originally a stage or a replacement, but wasn't
# replaced, should be in the final graph.
@@ -658,9 +695,9 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
data_side_input = {}
data_output = {}
for transform in stage.transforms:
- pcoll_id = transform.spec.payload
if transform.spec.urn in (bundle_processor.DATA_INPUT_URN,
bundle_processor.DATA_OUTPUT_URN):
+ pcoll_id = transform.spec.payload
if transform.spec.urn == bundle_processor.DATA_INPUT_URN:
target = transform.unique_name, only_element(transform.outputs)
data_input[target] = pcoll_id
@@ -673,13 +710,18 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
transform.spec.payload = data_operation_spec.SerializeToString()
else:
transform.spec.payload = ""
+ elif transform.spec.urn == urns.PARDO_TRANSFORM:
+ payload = proto_utils.parse_Bytes(
+ transform.spec.payload, beam_runner_api_pb2.ParDoPayload)
+ for tag, si in payload.side_inputs.items():
+ data_side_input[transform.unique_name, tag] = (
+ 'materialize:' + transform.inputs[tag],
+ beam.pvalue.SideInputData.from_runner_api(si, None))
return data_input, data_side_input, data_output
logging.info('Running %s', stage.name)
logging.debug(' %s', stage)
data_input, data_side_input, data_output = extract_endpoints(stage)
- if data_side_input:
- raise NotImplementedError('Side inputs.')
process_bundle_descriptor = beam_fn_api_pb2.ProcessBundleDescriptor(
id=self._next_uid(),
@@ -711,6 +753,20 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
data_out.write(element_data)
data_out.close()
+ # Store the required side inputs into state.
+ for (transform_id, tag), (pcoll_id, si) in data_side_input.items():
+ elements_by_window = _WindowGroupingBuffer(si)
+ for element_data in pcoll_buffers[pcoll_id]:
+ elements_by_window.append(element_data)
+ for window, elements_data in elements_by_window.items():
+ state_key = beam_fn_api_pb2.StateKey(
+ multimap_side_input=beam_fn_api_pb2.StateKey.MultimapSideInput(
+ ptransform_id=transform_id,
+ side_input_id=tag,
+ window=window))
+ controller.state_handler.blocking_append(
+ state_key, elements_data, process_bundle.instruction_id)
+
# Register and start running the bundle.
controller.control_handler.push(process_bundle_registration)
controller.control_handler.push(process_bundle)
@@ -975,34 +1031,56 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
# These classes are used to interact with the worker.
- class SimpleState(object): # TODO(robertwb): Inherit from GRPC servicer.
+ class StateServicer(beam_fn_api_pb2_grpc.BeamFnStateServicer):
def __init__(self):
- self._all = collections.defaultdict(list)
+ self._lock = threading.Lock()
+ self._state = collections.defaultdict(list)
- def Get(self, state_key):
- return beam_fn_api_pb2.Elements.Data(
- data=''.join(self._all[self._to_key(state_key)]))
+ def blocking_get(self, state_key, instruction_reference=None):
+ with self._lock:
+ return ''.join(self._state[self._to_key(state_key)])
- def Append(self, state_key, data):
- self._all[self._to_key(state_key)].extend(data)
+ def blocking_append(self, state_key, data, instruction_reference=None):
+ with self._lock:
+ self._state[self._to_key(state_key)].append(data)
- def Clear(self, state_key):
- try:
- del self._all[self._to_key(state_key)]
- except KeyError:
- pass
+ def blocking_clear(self, state_key, instruction_reference=None):
+ with self._lock:
+ del self._state[self._to_key(state_key)]
@staticmethod
def _to_key(state_key):
- return state_key.window, state_key.key
+ return state_key.SerializeToString()
+
+ class GrpcStateServicer(
+ StateServicer, beam_fn_api_pb2_grpc.BeamFnStateServicer):
+ def State(self, request_stream, context=None):
+ # Note that this eagerly mutates state, assuming any failures are fatal.
+ # Thus it is safe to ignore instruction_reference.
+ for request in request_stream:
+ if request.get:
+ yield beam_fn_api_pb2.StateResponse(
+ id=request.id,
+ get=beam_fn_api_pb2.StateGetResponse(
+ data=self.blocking_get(request.state_key)))
+ elif request.append:
+ self.blocking_append(request.state_key, request.append.data)
+ yield beam_fn_api_pb2.StateResponse(
+ id=request.id,
+ append=beam_fn_api_pb2.AppendResponse())
+ elif request.clear:
+ self.blocking_clear(request.state_key)
+ yield beam_fn_api_pb2.StateResponse(
+ id=request.id,
+ clear=beam_fn_api_pb2.ClearResponse())
class DirectController(object):
"""An in-memory controller for fn API control, state and data planes."""
def __init__(self):
self._responses = []
- self.state_handler = FnApiRunner.SimpleState()
+ self.state_handler = FnApiRunner.StateServicer()
self.control_handler = self
self.data_plane_handler = data_plane.InMemoryDataChannel()
self.worker = sdk_worker.SdkWorker(
@@ -1032,7 +1110,6 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
def __init__(self, sdk_harness_factory=None):
self.sdk_harness_factory = sdk_harness_factory
- self.state_handler = FnApiRunner.SimpleState()
self.control_server = grpc.server(
futures.ThreadPoolExecutor(max_workers=10))
self.control_port = self.control_server.add_insecure_port('[::]:0')
@@ -1049,6 +1126,12 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
beam_fn_api_pb2_grpc.add_BeamFnDataServicer_to_server(
self.data_plane_handler, self.data_server)
+ # TODO(robertwb): Is sharing the control channel fine? Alternatively,
+ # how should this be plumbed?
+ self.state_handler = FnApiRunner.GrpcStateServicer()
+ beam_fn_api_pb2_grpc.add_BeamFnStateServicer_to_server(
+ self.state_handler, self.control_server)
+
logging.info('starting control server on port %s', self.control_port)
logging.info('starting data server on port %s', self.data_port)
self.data_server.start()
@@ -1056,7 +1139,8 @@ class FnApiRunner(maptask_executor_runner.MapTaskExecutorRunner):
self.worker = (self.sdk_harness_factory or sdk_worker.SdkHarness)(
'localhost:%s' % self.control_port)
- self.worker_thread = threading.Thread(target=self.worker.run)
+ self.worker_thread = threading.Thread(
+ name='run_worker', target=self.worker.run)
logging.info('starting worker')
self.worker_thread.start()
http://git-wip-us.apache.org/repos/asf/beam/blob/38556b78/sdks/python/apache_beam/runners/portability/fn_api_runner_test.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/portability/fn_api_runner_test.py b/sdks/python/apache_beam/runners/portability/fn_api_runner_test.py
index 31f1b6f..ea9ed1a 100644
--- a/sdks/python/apache_beam/runners/portability/fn_api_runner_test.py
+++ b/sdks/python/apache_beam/runners/portability/fn_api_runner_test.py
@@ -24,6 +24,7 @@ from apache_beam.runners.portability import fn_api_runner
from apache_beam.runners.portability import maptask_executor_runner_test
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
+from apache_beam.transforms import window
try:
from apache_beam.runners.worker.statesampler import DEFAULT_SAMPLING_PERIOD_MS
@@ -47,12 +48,40 @@ class FnApiRunnerTest(
pass
def test_pardo_side_inputs(self):
- # TODO(BEAM-1348): Enable once side inputs are supported in fn API.
- pass
-
- def test_pardo_unfusable_side_inputs(self):
- # TODO(BEAM-1348): Enable once side inputs are supported in fn API.
- pass
+ def cross_product(elem, sides):
+ for side in sides:
+ yield elem, side
+ with self.create_pipeline() as p:
+ main = p | 'main' >> beam.Create(['a', 'b', 'c'])
+ side = p | 'side' >> beam.Create(['x', 'y'])
+ assert_that(main | beam.FlatMap(cross_product, beam.pvalue.AsList(side)),
+ equal_to([('a', 'x'), ('b', 'x'), ('c', 'x'),
+ ('a', 'y'), ('b', 'y'), ('c', 'y')]))
+
+ # Now with some windowing.
+ pcoll = p | beam.Create(range(10)) | beam.Map(
+ lambda t: window.TimestampedValue(t, t))
+ # Intentionally choosing non-aligned windows to highlight the transition.
+ main = pcoll | 'WindowMain' >> beam.WindowInto(window.FixedWindows(5))
+ side = pcoll | 'WindowSide' >> beam.WindowInto(window.FixedWindows(7))
+ res = main | beam.Map(lambda x, s: (x, sorted(s)),
+ beam.pvalue.AsList(side))
+ assert_that(
+ res,
+ equal_to([
+ # The window [0, 5) maps to the window [0, 7).
+ (0, range(7)),
+ (1, range(7)),
+ (2, range(7)),
+ (3, range(7)),
+ (4, range(7)),
+ # The window [5, 10) maps to the window [7, 14).
+ (5, range(7, 10)),
+ (6, range(7, 10)),
+ (7, range(7, 10)),
+ (8, range(7, 10)),
+ (9, range(7, 10))]),
+ label='windowed')
def test_assert_that(self):
# TODO: figure out a way for fn_api_runner to parse and raise the
http://git-wip-us.apache.org/repos/asf/beam/blob/38556b78/sdks/python/apache_beam/runners/worker/bundle_processor.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor.py b/sdks/python/apache_beam/runners/worker/bundle_processor.py
index 1049ae1..689eab7 100644
--- a/sdks/python/apache_beam/runners/worker/bundle_processor.py
+++ b/sdks/python/apache_beam/runners/worker/bundle_processor.py
@@ -39,6 +39,7 @@ from apache_beam.runners import pipeline_context
from apache_beam.runners.dataflow.native_io import iobase as native_iobase
from apache_beam.runners.worker import operation_specs
from apache_beam.runners.worker import operations
+from apache_beam.transforms import sideinputs
from apache_beam.utils import counters
from apache_beam.utils import proto_utils
from apache_beam.utils import urns
@@ -62,10 +63,6 @@ OLD_DATAFLOW_RUNNER_HARNESS_PARDO_URN = 'urn:beam:dofn:javasdk:0.1'
OLD_DATAFLOW_RUNNER_HARNESS_READ_URN = 'urn:org.apache.beam:source:java:0.1'
-def side_input_tag(transform_id, tag):
- return str("%d[%s][%s]" % (len(transform_id), transform_id, tag))
-
-
class RunnerIOOperation(operations.Operation):
"""Common baseclass for runner harness IO operations."""
@@ -162,6 +159,44 @@ class SideInputSource(native_iobase.NativeSource,
yield self._coder.get_impl().decode_from_stream(input_stream, True)
+class StateBackedSideInputMap(object):
+ def __init__(self, state_handler, transform_id, tag, side_input_data):
+ self._state_handler = state_handler
+ self._transform_id = transform_id
+ self._tag = tag
+ self._side_input_data = side_input_data
+ self._element_coder = side_input_data.coder.wrapped_value_coder
+ self._target_window_coder = side_input_data.coder.window_coder
+ # TODO(robertwb): Limit the cache size.
+ # TODO(robertwb): Cross-bundle caching respecting cache tokens.
+ self._cache = {}
+
+ def __getitem__(self, window):
+ target_window = self._side_input_data.window_mapping_fn(window)
+ if target_window not in self._cache:
+ state_key = beam_fn_api_pb2.StateKey(
+ multimap_side_input=beam_fn_api_pb2.StateKey.MultimapSideInput(
+ ptransform_id=self._transform_id,
+ side_input_id=self._tag,
+ window=self._target_window_coder.encode(target_window)))
+ element_coder_impl = self._element_coder.get_impl()
+ state_handler = self._state_handler
+
+ class AllElements(object):
+ def __iter__(self):
+ # TODO(robertwb): Support pagination.
+ input_stream = coder_impl.create_InputStream(
+ state_handler.blocking_get(state_key, None))
+ while input_stream.size() > 0:
+ yield element_coder_impl.decode_from_stream(input_stream, True)
+ self._cache[target_window] = self._side_input_data.view_fn(AllElements())
+ return self._cache[target_window]
+
+ def is_globally_windowed(self):
+ return (self._side_input_data.window_mapping_fn
+ == sideinputs._global_window_mapping_fn)
+
+
def memoize(func):
cache = {}
missing = object()
@@ -201,10 +236,17 @@ class BundleProcessor(object):
descriptor, self.data_channel_factory, self.counter_factory,
self.state_sampler, self.state_handler)
+ def is_side_input(transform_proto, tag):
+ if transform_proto.spec.urn == urns.PARDO_TRANSFORM:
+ return tag in proto_utils.parse_Bytes(
+ transform_proto.spec.payload,
+ beam_runner_api_pb2.ParDoPayload).side_inputs
+
pcoll_consumers = collections.defaultdict(list)
for transform_id, transform_proto in descriptor.transforms.items():
- for pcoll_id in transform_proto.inputs.values():
- pcoll_consumers[pcoll_id].append(transform_id)
+ 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):
@@ -412,7 +454,6 @@ def create(factory, transform_id, transform_proto, parameter, consumers):
@BeamTransformFactory.register_urn(
urns.READ_TRANSFORM, beam_runner_api_pb2.ReadPayload)
def create(factory, transform_id, transform_proto, parameter, consumers):
- # The Dataflow runner harness strips the base64 encoding.
source = iobase.SourceBase.from_runner_api(parameter.source, factory.context)
spec = operation_specs.WorkerRead(
iobase.SourceBundle(1.0, source, None, None),
@@ -428,17 +469,9 @@ def create(factory, transform_id, transform_proto, parameter, consumers):
@BeamTransformFactory.register_urn(OLD_DATAFLOW_RUNNER_HARNESS_PARDO_URN, None)
-def create(factory, transform_id, transform_proto, parameter, consumers):
- dofn_data = pickler.loads(parameter)
- if len(dofn_data) == 2:
- # Has side input data.
- serialized_fn, side_input_data = dofn_data
- else:
- # No side input data.
- serialized_fn, side_input_data = parameter, []
+def create(factory, transform_id, transform_proto, serialized_fn, consumers):
return _create_pardo_operation(
- factory, transform_id, transform_proto, consumers,
- serialized_fn, side_input_data)
+ factory, transform_id, transform_proto, consumers, serialized_fn)
@BeamTransformFactory.register_urn(
@@ -446,31 +479,26 @@ def create(factory, transform_id, transform_proto, parameter, consumers):
def create(factory, transform_id, transform_proto, parameter, consumers):
assert parameter.do_fn.spec.urn == urns.PICKLED_DO_FN_INFO
serialized_fn = parameter.do_fn.spec.payload
- dofn_data = pickler.loads(serialized_fn)
- if len(dofn_data) == 2:
- # Has side input data.
- serialized_fn, side_input_data = dofn_data
- else:
- # No side input data.
- side_input_data = []
return _create_pardo_operation(
factory, transform_id, transform_proto, consumers,
- serialized_fn, side_input_data)
+ serialized_fn, parameter.side_inputs)
def _create_pardo_operation(
factory, transform_id, transform_proto, consumers,
- serialized_fn, side_input_data):
- def create_side_input(tag, coder):
- # TODO(robertwb): Extract windows (and keys) out of element data.
- # TODO(robertwb): Extract state key from ParDoPayload.
- return operation_specs.WorkerSideInputSource(
- tag=tag,
- source=SideInputSource(
- factory.state_handler,
- beam_fn_api_pb2.StateKey.MultimapSideInput(
- key=side_input_tag(transform_id, tag)),
- coder=coder))
+ serialized_fn, side_inputs_proto=None):
+
+ if side_inputs_proto:
+ tagged_side_inputs = [
+ (tag, beam.pvalue.SideInputData.from_runner_api(si, factory.context))
+ for tag, si in side_inputs_proto.items()]
+ tagged_side_inputs.sort(key=lambda tag_si: int(tag_si[0][4:]))
+ side_input_maps = [
+ StateBackedSideInputMap(factory.state_handler, transform_id, tag, si)
+ for tag, si in tagged_side_inputs]
+ else:
+ side_input_maps = []
+
output_tags = list(transform_proto.outputs.keys())
# Hack to match out prefix injected by dataflow runner.
@@ -482,27 +510,31 @@ def _create_pardo_operation(
return 'out_' + tag
else:
return tag
+
dofn_data = pickler.loads(serialized_fn)
if not dofn_data[-1]:
# Windowing not set.
- pcoll_id, = transform_proto.inputs.values()
+ side_input_tags = side_inputs_proto or ()
+ pcoll_id, = [pcoll for tag, pcoll in transform_proto.inputs.items()
+ if tag not in side_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,))
+
output_coders = factory.get_output_coders(transform_proto)
spec = operation_specs.WorkerDoFn(
serialized_fn=serialized_fn,
output_tags=[mutate_tag(tag) for tag in output_tags],
input=None,
- side_inputs=[
- create_side_input(tag, coder) for tag, coder in side_input_data],
+ side_inputs=[], # Obsoleted by side_input_maps.
output_coders=[output_coders[tag] for tag in output_tags])
return factory.augment_oldstyle_op(
operations.DoOperation(
transform_proto.unique_name,
spec,
factory.counter_factory,
- factory.state_sampler),
+ factory.state_sampler,
+ side_input_maps),
transform_proto.unique_name,
consumers,
output_tags)
@@ -511,10 +543,8 @@ def _create_pardo_operation(
def _create_simple_pardo_operation(
factory, transform_id, transform_proto, consumers, dofn):
serialized_fn = pickler.dumps((dofn, (), {}, [], None))
- side_input_data = []
return _create_pardo_operation(
- factory, transform_id, transform_proto, consumers,
- serialized_fn, side_input_data)
+ factory, transform_id, transform_proto, consumers, serialized_fn)
@BeamTransformFactory.register_urn(
http://git-wip-us.apache.org/repos/asf/beam/blob/38556b78/sdks/python/apache_beam/runners/worker/data_plane.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/data_plane.py b/sdks/python/apache_beam/runners/worker/data_plane.py
index 5a511a0..f2a3751 100644
--- a/sdks/python/apache_beam/runners/worker/data_plane.py
+++ b/sdks/python/apache_beam/runners/worker/data_plane.py
@@ -25,6 +25,7 @@ import abc
import collections
import logging
import Queue as queue
+import sys
import threading
import grpc
@@ -147,6 +148,7 @@ class _GrpcDataChannel(DataChannel):
self._receive_lock = threading.Lock()
self._reads_finished = threading.Event()
self._closed = False
+ self._exc_info = None
def close(self):
self._to_send.put(self._WRITES_FINISHED)
@@ -163,12 +165,17 @@ class _GrpcDataChannel(DataChannel):
received = self._receiving_queue(instruction_id)
done_targets = []
while len(done_targets) < len(expected_targets):
- data = received.get()
- if not data.data and data.target in expected_targets:
- done_targets.append(data.target)
+ try:
+ data = received.get(timeout=1)
+ except queue.Empty:
+ if self._exc_info:
+ raise exc_info[0], exc_info[1], exc_info[2]
else:
- assert data.target not in done_targets
- yield data
+ if not data.data and data.target in expected_targets:
+ done_targets.append(data.target)
+ else:
+ assert data.target not in done_targets
+ yield data
def output_stream(self, instruction_id, target):
# TODO: Return an output stream that sends data
@@ -215,6 +222,7 @@ class _GrpcDataChannel(DataChannel):
except: # pylint: disable=bare-except
if not self._closed:
logging.exception('Failed to read inputs in the data plane')
+ self._exc_info = sys.exc_info()
raise
finally:
self._reads_finished.set()
http://git-wip-us.apache.org/repos/asf/beam/blob/38556b78/sdks/python/apache_beam/runners/worker/operations.pxd
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/operations.pxd b/sdks/python/apache_beam/runners/worker/operations.pxd
index d380a45..cb05c90 100644
--- a/sdks/python/apache_beam/runners/worker/operations.pxd
+++ b/sdks/python/apache_beam/runners/worker/operations.pxd
@@ -72,6 +72,7 @@ cdef class DoOperation(Operation):
cdef object dofn_runner
cdef Receiver dofn_receiver
cdef object tagged_receivers
+ cdef object side_input_maps
cdef class CombineOperation(Operation):
cdef object phased_combine_fn
http://git-wip-us.apache.org/repos/asf/beam/blob/38556b78/sdks/python/apache_beam/runners/worker/operations.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/operations.py b/sdks/python/apache_beam/runners/worker/operations.py
index ed3f3b8..6b5f024 100644
--- a/sdks/python/apache_beam/runners/worker/operations.py
+++ b/sdks/python/apache_beam/runners/worker/operations.py
@@ -259,6 +259,11 @@ class _TaggedReceivers(dict):
class DoOperation(Operation):
"""A Do operation that will execute a custom DoFn for each input element."""
+ def __init__(
+ self, name, spec, counter_factory, sampler, side_input_maps=None):
+ super(DoOperation, self).__init__(name, spec, counter_factory, sampler)
+ self.side_input_maps = side_input_maps
+
def _read_side_inputs(self, tags_and_types):
"""Generator reading side inputs in the order prescribed by tags_and_types.
@@ -273,6 +278,10 @@ class DoOperation(Operation):
either in singleton or collection mode according to the tags_and_types
argument.
"""
+ # Only call this on the old path where side_input_maps was not
+ # provided directly.
+ assert self.side_input_maps is None
+
# We will read the side inputs in the order prescribed by the
# tags_and_types argument because this is exactly the order needed to
# replace the ArgumentPlaceholder objects in the args/kwargs of the DoFn
@@ -336,8 +345,14 @@ class DoOperation(Operation):
raise ValueError('Unexpected output name for operation: %s' % tag)
self.tagged_receivers[original_tag] = self.receivers[index]
+ if self.side_input_maps is None:
+ if tags_and_types:
+ self.side_input_maps = list(self._read_side_inputs(tags_and_types))
+ else:
+ self.side_input_maps = []
+
self.dofn_runner = common.DoFnRunner(
- fn, args, kwargs, self._read_side_inputs(tags_and_types),
+ fn, args, kwargs, self.side_input_maps,
window_fn, context, self.tagged_receivers,
logger, self.step_name,
scoped_metrics_container=self.scoped_metrics_container)
http://git-wip-us.apache.org/repos/asf/beam/blob/38556b78/sdks/python/apache_beam/runners/worker/sdk_worker.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/runners/worker/sdk_worker.py b/sdks/python/apache_beam/runners/worker/sdk_worker.py
index b08e473..55ecbcc 100644
--- a/sdks/python/apache_beam/runners/worker/sdk_worker.py
+++ b/sdks/python/apache_beam/runners/worker/sdk_worker.py
@@ -24,6 +24,8 @@ from __future__ import print_function
import functools
import logging
import Queue as queue
+import sys
+import threading
import traceback
from concurrent import futures
@@ -46,9 +48,10 @@ class SdkHarness(object):
def run(self):
control_stub = beam_fn_api_pb2_grpc.BeamFnControlStub(self._control_channel)
- # TODO(robertwb): Wire up to new state api.
- state_stub = None
- self.worker = SdkWorker(state_stub, self._data_channel_factory)
+ state_stub = beam_fn_api_pb2_grpc.BeamFnStateStub(self._control_channel)
+ state_handler = GrpcStateHandler(state_stub)
+ state_handler.start()
+ self.worker = SdkWorker(state_handler, self._data_channel_factory)
responses = queue.Queue()
no_more_work = object()
@@ -102,6 +105,7 @@ class SdkHarness(object):
# control to its caller.
responses.put(no_more_work)
self._data_channel_factory.close()
+ state_handler.done()
logging.info('Done consuming work.')
@@ -148,3 +152,104 @@ class SdkWorker(object):
def process_bundle_progress(self, request, instruction_id):
# It is an error to get progress for a not-in-flight bundle.
return self.bundle_processors.get(instruction_id).metrics()
+
+
+class GrpcStateHandler(object):
+
+ _DONE = object()
+
+ def __init__(self, state_stub):
+ self._lock = threading.Lock()
+ self._state_stub = state_stub
+ self._requests = queue.Queue()
+ self._responses_by_id = {}
+ self._last_id = 0
+ self._exc_info = None
+
+ def start(self):
+ self._done = False
+
+ def request_iter():
+ while True:
+ request = self._requests.get()
+ if request is self._DONE or self._done:
+ break
+ yield request
+ responses = self._state_stub.State(request_iter())
+
+ def pull_responses():
+ try:
+ for response in responses:
+ self._responses_by_id[response.id].set(response)
+ if self._done:
+ break
+ except: # pylint: disable=bare-except
+ self._exc_info = sys.exc_info()
+ raise
+ reader = threading.Thread(target=pull_responses, name='read_state')
+ reader.daemon = True
+ reader.start()
+
+ def done(self):
+ self._done = True
+ self._requests.put(self._DONE)
+
+ def blocking_get(self, state_key, instruction_reference):
+ response = self._blocking_request(
+ beam_fn_api_pb2.StateRequest(
+ instruction_reference=instruction_reference,
+ state_key=state_key,
+ get=beam_fn_api_pb2.StateGetRequest()))
+ if response.get.continuation_token:
+ raise NotImplementedErrror
+ return response.get.data
+
+ def blocking_append(self, state_key, data, instruction_reference):
+ self._blocking_request(
+ beam_fn_api_pb2.StateRequest(
+ instruction_reference=instruction_reference,
+ state_key=state_key,
+ append=beam_fn_api_pb2.StateAppendRequest(data=data)))
+
+ def blocking_clear(self, state_key, instruction_reference):
+ self._blocking_request(
+ beam_fn_api_pb2.StateRequest(
+ instruction_reference=instruction_reference,
+ state_key=state_key,
+ clear=beam_fn_api_pb2.StateClearRequest()))
+
+ def _blocking_request(self, request):
+ request.id = self._next_id()
+ self._responses_by_id[request.id] = future = _Future()
+ self._requests.put(request)
+ while not future.wait(timeout=1):
+ if self._exc_info:
+ raise self._exc_info[0], self._exc_info[1], self._exc_info[2]
+ elif self._done:
+ raise RuntimeError()
+ del self._responses_by_id[request.id]
+ return future.get()
+
+ def _next_id(self):
+ self._last_id += 1
+ return str(self._last_id)
+
+
+class _Future(object):
+ """A simple future object to implement blocking requests.
+ """
+ def __init__(self):
+ self._event = threading.Event()
+
+ def wait(self, timeout=None):
+ return self._event.wait(timeout)
+
+ def get(self, timeout=None):
+ if self.wait(timeout):
+ return self._value
+ else:
+ raise LookupError()
+
+ def set(self, value):
+ self._value = value
+ self._event.set()
http://git-wip-us.apache.org/repos/asf/beam/blob/38556b78/sdks/python/apache_beam/transforms/core.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/transforms/core.py b/sdks/python/apache_beam/transforms/core.py
index ff2428e..fbc5f33 100644
--- a/sdks/python/apache_beam/transforms/core.py
+++ b/sdks/python/apache_beam/transforms/core.py
@@ -842,7 +842,7 @@ class ParDo(PTransformWithSideInputs):
return _MultiParDo(self, tags, main_tag)
def _pardo_fn_data(self):
- si_tags_and_types = []
+ si_tags_and_types = None
windowing = None
return self.fn, self.args, self.kwargs, si_tags_and_types, windowing
@@ -855,7 +855,15 @@ class ParDo(PTransformWithSideInputs):
do_fn=beam_runner_api_pb2.SdkFunctionSpec(
spec=beam_runner_api_pb2.FunctionSpec(
urn=urns.PICKLED_DO_FN_INFO,
- payload=picked_pardo_fn_data))))
+ payload=picked_pardo_fn_data)),
+ # It'd be nice to name these according to their actual
+ # names/positions in the orignal argument list, but such a
+ # transformation is currently irreversible given how
+ # remove_objects_from_args and insert_values_in_args
+ # are currently implemented.
+ side_inputs={
+ "side%s" % ix: si.to_runner_api(context)
+ for ix, si in enumerate(self.side_inputs)}))
@PTransform.register_urn(
urns.PARDO_TRANSFORM, beam_runner_api_pb2.ParDoPayload)
@@ -864,10 +872,17 @@ class ParDo(PTransformWithSideInputs):
fn, args, kwargs, si_tags_and_types, windowing = pickler.loads(
pardo_payload.do_fn.spec.payload)
if si_tags_and_types:
- raise NotImplementedError('deferred side inputs')
+ raise NotImplementedError('explicit side input data')
elif windowing:
raise NotImplementedError('explicit windowing')
- return ParDo(fn, *args, **kwargs)
+ result = ParDo(fn, *args, **kwargs)
+ # This is an ordered list stored as a dict (see the comments in
+ # to_runner_api_parameter above).
+ indexed_side_inputs = [
+ (int(ix[4:]), pvalue.AsSideInput.from_runner_api(si, context))
+ for ix, si in pardo_payload.side_inputs.items()]
+ result.side_inputs = [si for _, si in sorted(indexed_side_inputs)]
+ return result
class _MultiParDo(PTransform):
http://git-wip-us.apache.org/repos/asf/beam/blob/38556b78/sdks/python/apache_beam/utils/urns.py
----------------------------------------------------------------------
diff --git a/sdks/python/apache_beam/utils/urns.py b/sdks/python/apache_beam/utils/urns.py
index 2aeaa53..c6135ba 100644
--- a/sdks/python/apache_beam/utils/urns.py
+++ b/sdks/python/apache_beam/utils/urns.py
@@ -58,6 +58,11 @@ LENGTH_PREFIX_CODER = "urn:beam:coders:length_prefix:0.1"
GLOBAL_WINDOW_CODER = "urn:beam:coders:urn:beam:coders:global_window:0.1"
WINDOWED_VALUE_CODER = "urn:beam:coders:windowed_value:0.1"
+ITERABLE_ACCESS = "urn:beam:sideinput:iterable"
+MULTIMAP_ACCESS = "urn:beam:sideinput:multimap"
+PICKLED_PYTHON_VIEWFN = "beam:view_fn:pickled_python_data:v0.1"
+PICKLED_WINDOW_MAPPING_FN = "beam:window_mapping_fn:pickled_python:v0.1"
+
class RunnerApiFn(object):
"""Abstract base class that provides urn registration utilities.