Move the pipeline graph manipulation utilites to their own files.

As well as being very different from the low-level execution details
found in common.py, there is no need to compile these (whereas the
other structures in common.py are critical to optimize for elementwise
efficiency).

Author: robertwb

sdks/python/apache_beam/pipeline.py

@@ -87,8 +87,8 @@
 from apache_beam.portability import common_urns
 from apache_beam.portability.api import beam_runner_api_pb2
 from apache_beam.runners import PipelineRunner
-from apache_beam.runners import common
 from apache_beam.runners import create_runner
+from apache_beam.runners import pipeline_utils
 from apache_beam.transforms import ParDo
 from apache_beam.transforms import ptransform
 from apache_beam.transforms.display import DisplayData
@@ -1019,7 +1019,7 @@ def merge_compatible_environments(proto):
 
     Mutates proto as contexts may have references to proto.components.
     """
-    common.merge_common_environments(proto, inplace=True)
+    pipeline_utils.merge_common_environments(proto, inplace=True)
 
   @staticmethod
   def from_runner_api(

sdks/python/apache_beam/runners/common.py

@@ -1920,230 +1920,3 @@ def windows(self):
       raise AttributeError('windows not accessible in this context')
     else:
       return self.windowed_value.windows
-
-
-def group_by_key_input_visitor(deterministic_key_coders=True):
-  # Importing here to avoid a circular dependency
-  # pylint: disable=wrong-import-order, wrong-import-position
-  from apache_beam.pipeline import PipelineVisitor
-  from apache_beam.transforms.core import GroupByKey
-
-  class GroupByKeyInputVisitor(PipelineVisitor):
-    """A visitor that replaces `Any` element type for input `PCollection` of
-    a `GroupByKey` with a `KV` type.
-
-    TODO(BEAM-115): Once Python SDK is compatible with the new Runner API,
-    we could directly replace the coder instead of mutating the element type.
-    """
-    def __init__(self, deterministic_key_coders=True):
-      self.deterministic_key_coders = deterministic_key_coders
-
-    def enter_composite_transform(self, transform_node):
-      self.visit_transform(transform_node)
-
-    def visit_transform(self, transform_node):
-      if isinstance(transform_node.transform, GroupByKey):
-        pcoll = transform_node.inputs[0]
-        pcoll.element_type = typehints.coerce_to_kv_type(
-            pcoll.element_type, transform_node.full_label)
-        pcoll.requires_deterministic_key_coder = (
-            self.deterministic_key_coders and transform_node.full_label)
-        key_type, value_type = pcoll.element_type.tuple_types
-        if transform_node.outputs:
-          key = next(iter(transform_node.outputs.keys()))
-          transform_node.outputs[key].element_type = typehints.KV[
-              key_type, typehints.Iterable[value_type]]
-          transform_node.outputs[key].requires_deterministic_key_coder = (
-              self.deterministic_key_coders and transform_node.full_label)
-
-  return GroupByKeyInputVisitor(deterministic_key_coders)
-
-
-def validate_pipeline_graph(pipeline_proto):
-  """Ensures this is a correctly constructed Beam pipeline.
-  """
-  def get_coder(pcoll_id):
-    return pipeline_proto.components.coders[
-        pipeline_proto.components.pcollections[pcoll_id].coder_id]
-
-  def validate_transform(transform_id):
-    transform_proto = pipeline_proto.components.transforms[transform_id]
-
-    # Currently the only validation we perform is that GBK operations have
-    # their coders set properly.
-    if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:
-      if len(transform_proto.inputs) != 1:
-        raise ValueError("Unexpected number of inputs: %s" % transform_proto)
-      if len(transform_proto.outputs) != 1:
-        raise ValueError("Unexpected number of outputs: %s" % transform_proto)
-      input_coder = get_coder(next(iter(transform_proto.inputs.values())))
-      output_coder = get_coder(next(iter(transform_proto.outputs.values())))
-      if input_coder.spec.urn != common_urns.coders.KV.urn:
-        raise ValueError(
-            "Bad coder for input of %s: %s" % (transform_id, input_coder))
-      if output_coder.spec.urn != common_urns.coders.KV.urn:
-        raise ValueError(
-            "Bad coder for output of %s: %s" % (transform_id, output_coder))
-      output_values_coder = pipeline_proto.components.coders[
-          output_coder.component_coder_ids[1]]
-      if (input_coder.component_coder_ids[0] !=
-          output_coder.component_coder_ids[0] or
-          output_values_coder.spec.urn != common_urns.coders.ITERABLE.urn or
-          output_values_coder.component_coder_ids[0] !=
-          input_coder.component_coder_ids[1]):
-        raise ValueError(
-            "Incompatible input coder %s and output coder %s for transform %s" %
-            (transform_id, input_coder, output_coder))
-    elif transform_proto.spec.urn == common_urns.primitives.ASSIGN_WINDOWS.urn:
-      if not transform_proto.inputs:
-        raise ValueError("Missing input for transform: %s" % transform_proto)
-    elif transform_proto.spec.urn == common_urns.primitives.PAR_DO.urn:
-      if not transform_proto.inputs:
-        raise ValueError("Missing input for transform: %s" % transform_proto)
-
-    for t in transform_proto.subtransforms:
-      validate_transform(t)
-
-  for t in pipeline_proto.root_transform_ids:
-    validate_transform(t)
-
-
-def _dep_key(dep):
-  if dep.type_urn == common_urns.artifact_types.FILE.urn:
-    payload = beam_runner_api_pb2.ArtifactFilePayload.FromString(
-        dep.type_payload)
-    if payload.sha256:
-      type_info = 'sha256', payload.sha256
-    else:
-      type_info = 'path', payload.path
-  elif dep.type_urn == common_urns.artifact_types.URL.urn:
-    payload = beam_runner_api_pb2.ArtifactUrlPayload.FromString(
-        dep.type_payload)
-    if payload.sha256:
-      type_info = 'sha256', payload.sha256
-    else:
-      type_info = 'url', payload.url
-  else:
-    type_info = dep.type_urn, dep.type_payload
-  return type_info, dep.role_urn, dep.role_payload
-
-
-def _expanded_dep_keys(dep):
-  if (dep.type_urn == common_urns.artifact_types.FILE.urn and
-      dep.role_urn == common_urns.artifact_roles.STAGING_TO.urn):
-    payload = beam_runner_api_pb2.ArtifactFilePayload.FromString(
-        dep.type_payload)
-    role = beam_runner_api_pb2.ArtifactStagingToRolePayload.FromString(
-        dep.role_payload)
-    if role.staged_name == 'submission_environment_dependencies.txt':
-      return
-    elif role.staged_name == 'requirements.txt':
-      with open(payload.path) as fin:
-        for line in fin:
-          yield 'requirements.txt', line.strip()
-      return
-
-  yield _dep_key(dep)
-
-
-def _base_env_key(env, include_deps=True):
-  return (
-      env.urn,
-      env.payload,
-      tuple(sorted(env.capabilities)),
-      tuple(sorted(env.resource_hints.items())),
-      tuple(sorted(_dep_key(dep)
-                   for dep in env.dependencies)) if include_deps else None)
-
-
-def _env_key(env):
-  return tuple(
-      sorted(
-          _base_env_key(e)
-          for e in environments.expand_anyof_environments(env)))
-
-
-def merge_common_environments(pipeline_proto, inplace=False):
-  canonical_environments = collections.defaultdict(list)
-  for env_id, env in pipeline_proto.components.environments.items():
-    canonical_environments[_env_key(env)].append(env_id)
-
-  if len(canonical_environments) == len(pipeline_proto.components.environments):
-    # All environments are already sufficiently distinct.
-    return pipeline_proto
-
-  environment_remappings = {
-      e: es[0]
-      for es in canonical_environments.values() for e in es
-  }
-
-  return update_environments(pipeline_proto, environment_remappings, inplace)
-
-
-def merge_superset_dep_environments(pipeline_proto):
-  """Merges all environemnts A and B where A and B are equivalent except that
-  A has a superset of the dependencies of B.
-  """
-  docker_envs = {}
-  for env_id, env in pipeline_proto.components.environments.items():
-    docker_env = environments.resolve_anyof_environment(
-        env, common_urns.environments.DOCKER.urn)
-    if docker_env.urn == common_urns.environments.DOCKER.urn:
-      docker_envs[env_id] = docker_env
-
-  has_base_and_dep = collections.defaultdict(set)
-  env_scores = {
-      env_id: (len(env.dependencies), env_id)
-      for (env_id, env) in docker_envs.items()
-  }
-
-  for env_id, env in docker_envs.items():
-    base_key = _base_env_key(env, include_deps=False)
-    has_base_and_dep[base_key, None].add(env_id)
-    for dep in env.dependencies:
-      for dep_key in _expanded_dep_keys(dep):
-        has_base_and_dep[base_key, dep_key].add(env_id)
-
-  environment_remappings = {}
-  for env_id, env in docker_envs.items():
-    base_key = _base_env_key(env, include_deps=False)
-    # This is the set of all environments that have at least all of env's deps.
-    candidates = set.intersection(
-        has_base_and_dep[base_key, None],
-        *[
-            has_base_and_dep[base_key, dep_key] for dep in env.dependencies
-            for dep_key in _expanded_dep_keys(dep)
-        ])
-    # Choose the maximal one.
-    best = max(candidates, key=env_scores.get)
-    if best != env_id:
-      environment_remappings[env_id] = best
-
-  return update_environments(pipeline_proto, environment_remappings)
-
-
-def update_environments(pipeline_proto, environment_remappings, inplace=False):
-  if not environment_remappings:
-    return pipeline_proto
-
-  if not inplace:
-    pipeline_proto = copy.copy(pipeline_proto)
-
-  for t in pipeline_proto.components.transforms.values():
-    if t.environment_id not in pipeline_proto.components.environments:
-      # TODO(https://github.com/apache/beam/issues/30876): Remove this
-      #  workaround.
-      continue
-    if t.environment_id and t.environment_id in environment_remappings:
-      t.environment_id = environment_remappings[t.environment_id]
-  for w in pipeline_proto.components.windowing_strategies.values():
-    if w.environment_id not in pipeline_proto.components.environments:
-      # TODO(https://github.com/apache/beam/issues/30876): Remove this
-      #  workaround.
-      continue
-    if w.environment_id and w.environment_id in environment_remappings:
-      w.environment_id = environment_remappings[w.environment_id]
-  for e in set(environment_remappings.keys()) - set(
-      environment_remappings.values()):
-    del pipeline_proto.components.environments[e]
-  return pipeline_proto