#

# 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.

#

"""Pipeline manipulation utilities useful for many runners.

For internal use only; no backwards-compatibility guarantees.

"""

# pytype: skip-fileimport collections

import collections

import copy

from apache_beam.portability import common_urns

from apache_beam.portability.api import beam_runner_api_pb2

from apache_beam.transforms import environments

from apache_beam.typehints import typehints

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