#

# 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, the top-level Beam object.

A pipeline holds a DAG of data transforms. Conceptually the nodes of the DAG

are transforms (:class:`~apache_beam.transforms.ptransform.PTransform` objects)

and the edges are values (mostly :class:`~apache_beam.pvalue.PCollection`

objects). The transforms take as inputs one or more PValues and output one or

more :class:`~apache_beam.pvalue.PValue` s.

The pipeline offers functionality to traverse the graph. The actual operation

to be executed for each node visited is specified through a runner object.

Typical usage::

# Create a pipeline object using a local runner for execution.

with beam.Pipeline('DirectRunner') as p:

# Add to the pipeline a "Create" transform. When executed this

# transform will produce a PCollection object with the specified values.

pcoll = p | 'Create' >> beam.Create([1, 2, 3])

# Another transform could be applied to pcoll, e.g., writing to a text file.

# For other transforms, refer to transforms/ directory.

pcoll | 'Write' >> beam.io.WriteToText('./output')

# run() will execute the DAG stored in the pipeline. The execution of the

# nodes visited is done using the specified local runner.

"""

# pytype: skip-file

# mypy: disallow-untyped-defs

import abc

import contextlib

import copy

import logging

import os

import re

import shutil

import tempfile

import threading

import unicodedata

import uuid

from collections import defaultdict

from collections.abc import Iterable

from collections.abc import Mapping

from collections.abc import Sequence

from typing import TYPE_CHECKING

from typing import Any

from typing import Optional

from typing import Type

from typing import Union

from google.protobuf import message

from apache_beam import pvalue

from apache_beam.coders import typecoders

from apache_beam.internal import pickler

from apache_beam.io.filesystems import FileSystems

from apache_beam.options.pipeline_options import CrossLanguageOptions

from apache_beam.options.pipeline_options import DebugOptions

from apache_beam.options.pipeline_options import PipelineOptions

from apache_beam.options.pipeline_options import SetupOptions

from apache_beam.options.pipeline_options import StandardOptions

from apache_beam.options.pipeline_options import StreamingOptions

from apache_beam.options.pipeline_options import TypeOptions

from apache_beam.options.pipeline_options_context import scoped_pipeline_options

from apache_beam.options.pipeline_options_validator import PipelineOptionsValidator

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

from apache_beam.transforms.display import HasDisplayData

from apache_beam.transforms.resources import merge_resource_hints

from apache_beam.transforms.resources import resource_hints_from_options

from apache_beam.transforms.sideinputs import get_sideinput_index

from apache_beam.typehints import TypeCheckError

from apache_beam.typehints import typehints

from apache_beam.utils import proto_utils

from apache_beam.utils import subprocess_server

from apache_beam.utils.annotations import deprecated

from apache_beam.utils.interactive_utils import alter_label_if_ipython

from apache_beam.utils.interactive_utils import is_in_ipython

if TYPE_CHECKING:

from types import TracebackType

from apache_beam.runners.pipeline_context import PipelineContext

from apache_beam.runners.runner import PipelineResult

from apache_beam.transforms import environments

__all__ = ['Pipeline', 'transform_annotations']

class Pipeline(HasDisplayData):

"""A pipeline object that manages a DAG of

:class:`~apache_beam.transforms.ptransform.PTransform` s

and their :class:`~apache_beam.pvalue.PValue` s.

Conceptually the :class:`~apache_beam.transforms.ptransform.PTransform` s are

the DAG's nodes and the :class:`~apache_beam.pvalue.PValue` s are the edges.

All the transforms applied to the pipeline must have distinct full labels.

If same transform instance needs to be applied then the right shift operator

should be used to designate new names

(e.g. ``input | "label" >> my_transform``).

"""

@classmethod

def runner_implemented_transforms(cls) -> frozenset[str]:

# This set should only contain transforms which are required to be

# implemented by a runner.

return frozenset([

common_urns.primitives.GROUP_BY_KEY.urn,

common_urns.primitives.IMPULSE.urn,

])

def __init__(

self,

runner: Optional[Union[str, PipelineRunner]] = None,

options: Optional[PipelineOptions] = None,

argv: Optional[list[str]] = None,

display_data: Optional[dict[str, Any]] = None):

"""Initialize a pipeline object.

Args:

runner (~apache_beam.runners.runner.PipelineRunner): An object of

type :class:`~apache_beam.runners.runner.PipelineRunner` that will be

used to execute the pipeline. For registered runners, the runner name

can be specified, otherwise a runner object must be supplied.

options (~apache_beam.options.pipeline_options.PipelineOptions):

A configured

:class:`~apache_beam.options.pipeline_options.PipelineOptions` object

containing arguments that should be used for running the Beam job.

argv (list[str]): a list of arguments (such as :data:`sys.argv`)

to be used for building a

:class:`~apache_beam.options.pipeline_options.PipelineOptions` object.

This will only be used if argument **options** is :data:`None`.

display_data (dict[str, Any]): a dictionary of static data associated

with this pipeline that can be displayed when it runs.

Raises:

ValueError: if either the runner or options argument is not

of the expected type.

"""

# Initializing logging configuration in case the user did not set it up.

logging.basicConfig()

if options is not None:

if isinstance(options, PipelineOptions):

# Make a deep copy of options since they could be overwritten in later

# steps. However, the 'runner' object within 'options' is excluded from

# the deep copy (it is shallow copied) due to potential issues with deep

# copying specific runner instances, such as FlumeRunner.

saved_runner = options.view_as(StandardOptions).runner

options.view_as(StandardOptions).runner = None

self._options = copy.deepcopy(options)

self._options.view_as(StandardOptions).runner = saved_runner

options.view_as(StandardOptions).runner = saved_runner

else:

raise ValueError(

'Parameter options, if specified, must be of type PipelineOptions. '

'Received : %r' % options)

elif argv is not None:

if isinstance(argv, list):

self._options = PipelineOptions(argv)

else:

raise ValueError(

'Parameter argv, if specified, must be a list. Received : %r' %

argv)

else:

self._options = PipelineOptions([])

FileSystems.set_options(self._options)

if runner is None:

runner = self._options.view_as(StandardOptions).runner

if runner is None:

runner = StandardOptions.DEFAULT_RUNNER

logging.info((

'Missing pipeline option (runner). Executing pipeline '

'using the default runner: %s.'),

runner)

if isinstance(runner, str):

runner = create_runner(runner)

elif not isinstance(runner, PipelineRunner):

raise TypeError(

'Runner %s is not a PipelineRunner object or the '

'name of a registered runner.' % runner)

# Runner can override the default pickler to be used.

if (self._options.view_as(SetupOptions).pickle_library == 'default' and

runner.default_pickle_library_override()):

logging.info(

"Runner defaulting to pickling library: %s.",

runner.default_pickle_library_override())

self._options.view_as(

SetupOptions).pickle_library = runner.default_pickle_library_override(

)

pickler.set_library(self._options.view_as(SetupOptions).pickle_library)

# Validate pipeline options

errors = PipelineOptionsValidator(self._options, runner).validate()

if errors:

raise ValueError(

'Pipeline has validations errors: \n' + '\n'.join(errors))

typecoders.registry.update_compatibility_version = self._options.view_as(

StreamingOptions).update_compatibility_version

# set default experiments for portable runners

# (needs to occur prior to pipeline construction)

if runner.is_fnapi_compatible():

experiments = (self._options.view_as(DebugOptions).experiments or [])

if not 'beam_fn_api' in experiments:

experiments.append('beam_fn_api')

self._options.view_as(DebugOptions).experiments = experiments

self.local_tempdir = tempfile.mkdtemp(prefix='beam-pipeline-temp')

# Default runner to be used.

self.runner = runner

# Stack of transforms generated by nested apply() calls. The stack will

# contain a root node as an enclosing (parent) node for top transforms.

self.transforms_stack = [

AppliedPTransform(None, None, '', None, None, None)

]

# Set of transform labels (full labels) applied to the pipeline.

# If a transform is applied and the full label is already in the set

# then the transform will have to be cloned with a new label.

self.applied_labels: set[str] = set()

# Hints supplied via pipeline options are considered the outermost hints.

self._root_transform().resource_hints = resource_hints_from_options(options)

# Create a ComponentIdMap for assigning IDs to components. Ensures that any

# components that receive an ID during pipeline construction (for example in

# ExternalTransform), will receive the same component ID when generating the

# full pipeline proto.

self.component_id_map = ComponentIdMap()

# Records whether this pipeline contains any external transforms.

self.contains_external_transforms = False

self._display_data = display_data or {}

self._error_handlers = []

self._annotations_stack = [{}]

def display_data(self) -> dict[str, Any]:

return self._display_data

@property # type: ignore[misc] # decorated property not supported

def options(self) -> PipelineOptions:

return self._options

@property

def allow_unsafe_triggers(self) -> bool:

return self._options.view_as(TypeOptions).allow_unsafe_triggers

def _register_error_handler(self, error_handler):

self._error_handlers.append(error_handler)

def _current_transform(self) -> 'AppliedPTransform':

"""Returns the transform currently on the top of the stack."""

return self.transforms_stack[-1]

@contextlib.contextmanager

def transform_annotations(self, **annotations):

"""A context manager for attaching annotations to a set of transforms.

All transforms applied while this context is active will have these

annotations attached. This includes sub-transforms applied within

composite transforms.

"""

self._annotations_stack.append({

**self._annotations_stack[-1], **encode_annotations(annotations)

})

yield

self._annotations_stack.pop()

def _current_annotations(self):

"""Returns the set of annotations that should be used on apply."""

return {**_global_annotations_stack()[-1], **self._annotations_stack[-1]}

def _root_transform(self) -> 'AppliedPTransform':

"""Returns the root transform of the transform stack."""

return self.transforms_stack[0]

def _remove_labels_recursively(

self, applied_transform: 'AppliedPTransform') -> None:

for part in applied_transform.parts:

if part.full_label in self.applied_labels:

self.applied_labels.remove(part.full_label)

self._remove_labels_recursively(part)

def _replace(self, override: 'PTransformOverride') -> None:

assert isinstance(override, PTransformOverride)

# From original transform output --> replacement transform output

output_map: dict[pvalue.PValue, pvalue.PValue] = {}

output_replacements: dict[AppliedPTransform,

list[tuple[pvalue.PValue, Optional[str]]]] = {}

input_replacements: dict[AppliedPTransform,

Mapping[str,

Union[pvalue.PBegin,

pvalue.PCollection]]] = {}

side_input_replacements: dict[AppliedPTransform,

list[pvalue.AsSideInput]] = {}

class TransformUpdater(PipelineVisitor): # pylint: disable=used-before-assignment

""""A visitor that replaces the matching PTransforms."""

def __init__(self, pipeline: Pipeline) -> None:

self.pipeline = pipeline

def _replace_if_needed(

self, original_transform_node: AppliedPTransform) -> None:

if override.matches(original_transform_node):

assert isinstance(original_transform_node, AppliedPTransform)

replacement_transform = (

override.get_replacement_transform_for_applied_ptransform(

original_transform_node))

if replacement_transform is original_transform_node.transform:

return

replacement_transform.side_inputs = tuple(

getattr(original_transform_node.transform, 'side_inputs', ()))

replacement_transform_node = AppliedPTransform(

original_transform_node.parent,

replacement_transform,

original_transform_node.full_label,

original_transform_node.main_inputs,

None,

annotations=original_transform_node.annotations)

# TODO(https://github.com/apache/beam/issues/21178): Merge rather

# than override.

replacement_transform_node.resource_hints = (

original_transform_node.resource_hints)

# Transform execution could depend on order in which nodes are

# considered. Hence we insert the replacement transform node to same

# index as the original transform node. Note that this operation

# removes the original transform node.

if original_transform_node.parent:

assert isinstance(original_transform_node.parent, AppliedPTransform)

parent_parts = original_transform_node.parent.parts

parent_parts[parent_parts.index(original_transform_node)] = (

replacement_transform_node)

else:

# Original transform has to be a root.

roots = self.pipeline.transforms_stack[0].parts

assert original_transform_node in roots

roots[roots.index(original_transform_node)] = (

replacement_transform_node)

inputs = override.get_replacement_inputs(original_transform_node)

if len(inputs) > 1:

transform_input = inputs

elif len(inputs) == 1:

transform_input = inputs[0]

elif len(inputs) == 0:

transform_input = pvalue.PBegin(self.pipeline)

try:

# We have to add the new AppliedTransform to the stack before

# expand() and pop it out later to make sure that parts get added

# correctly.

self.pipeline.transforms_stack.append(replacement_transform_node)

# Keeping the same label for the replaced node but recursively

# removing labels of child transforms of original transform since

# they will be replaced during the expand below. This is needed in

# case the replacement contains children that have labels that

# conflicts with labels of the children of the original.

self.pipeline._remove_labels_recursively(original_transform_node)

new_output = replacement_transform.expand(transform_input)

assert isinstance(

new_output, (dict, pvalue.PValue, pvalue.DoOutputsTuple))

if isinstance(new_output, pvalue.PValue):

new_output.element_type = None

self.pipeline._infer_result_type(

replacement_transform, inputs, new_output)

if isinstance(new_output, dict):

for new_tag, new_pcoll in new_output.items():

replacement_transform_node.add_output(new_pcoll, new_tag)

elif isinstance(new_output, pvalue.DoOutputsTuple):

replacement_transform_node.add_output(

new_output, new_output._main_tag)

else:

replacement_transform_node.add_output(new_output, new_output.tag)

# Recording updated outputs. This cannot be done in the same

# visitor since if we dynamically update output type here, we'll

# run into errors when visiting child nodes.

#

# NOTE: When replacing multiple outputs, the replacement

# PCollection tags must have a matching tag in the original

# transform.

if isinstance(new_output, pvalue.PValue):

if not new_output.producer:

new_output.producer = replacement_transform_node

output_map[original_transform_node.outputs[new_output.tag]] = \

new_output

elif isinstance(new_output, (pvalue.DoOutputsTuple, tuple)):

for pcoll in new_output:

if not pcoll.producer:

pcoll.producer = replacement_transform_node

output_map[original_transform_node.outputs[pcoll.tag]] = pcoll

elif isinstance(new_output, dict):

for tag, pcoll in new_output.items():

if not pcoll.producer:

pcoll.producer = replacement_transform_node

output_map[original_transform_node.outputs[tag]] = pcoll

finally:

self.pipeline.transforms_stack.pop()

def enter_composite_transform(

self, transform_node: AppliedPTransform) -> None:

self._replace_if_needed(transform_node)

def visit_transform(self, transform_node: AppliedPTransform) -> None:

self._replace_if_needed(transform_node)

self.visit(TransformUpdater(self))

# Ensure no type information is lost.

for old, new in output_map.items():

if new.element_type == typehints.Any:

# TODO(robertwb): Perhaps take the intersection?

new.element_type = old.element_type

# Adjusting inputs and outputs

class InputOutputUpdater(PipelineVisitor): # pylint: disable=used-before-assignment

""""A visitor that records input and output values to be replaced.

Input and output values that should be updated are recorded in maps

input_replacements and output_replacements respectively.

We cannot update input and output values while visiting since that results

in validation errors.

"""

def __init__(self, pipeline: Pipeline) -> None:

self.pipeline = pipeline

def enter_composite_transform(

self, transform_node: AppliedPTransform) -> None:

self.visit_transform(transform_node)

def visit_transform(self, transform_node: AppliedPTransform) -> None:

replace_output = False

for tag in transform_node.outputs:

if transform_node.outputs[tag] in output_map:

replace_output = True

break

replace_input = False

for input in transform_node.inputs:

if input in output_map:

replace_input = True

break

replace_side_inputs = False

for side_input in transform_node.side_inputs:

if side_input.pvalue in output_map:

replace_side_inputs = True

break

if replace_output:

output_replacements[transform_node] = []

for original, replacement in output_map.items():

for tag, output in transform_node.outputs.items():

if output == original:

output_replacements[transform_node].append((tag, replacement))

if replace_input:

new_inputs = {

tag: input if not input in output_map else output_map[input]

for (tag, input) in transform_node.main_inputs.items()

}

input_replacements[transform_node] = new_inputs

if replace_side_inputs:

new_side_inputs = []

for side_input in transform_node.side_inputs:

if side_input.pvalue in output_map:

side_input.pvalue = output_map[side_input.pvalue]

new_side_inputs.append(side_input)

else:

new_side_inputs.append(side_input)

side_input_replacements[transform_node] = new_side_inputs

self.visit(InputOutputUpdater(self))

for transform, output_replacement in output_replacements.items():

for tag, output in output_replacement:

transform.replace_output(output, tag=tag)

for transform, input_replacement in input_replacements.items():

transform.replace_inputs(input_replacement)

for transform, side_input_replacement in side_input_replacements.items():

transform.replace_side_inputs(side_input_replacement)

def _check_replacement(self, override: 'PTransformOverride') -> None:

class ReplacementValidator(PipelineVisitor):

def visit_transform(self, transform_node: AppliedPTransform) -> None:

if override.matches(transform_node):

raise RuntimeError(

'Transform node %r was not replaced as expected.' %

transform_node)

self.visit(ReplacementValidator())

def replace_all(self, replacements: Iterable['PTransformOverride']) -> None:

""" Dynamically replaces PTransforms in the currently populated hierarchy.

Currently this only works for replacements where input and output types

are exactly the same.

TODO: Update this to also work for transform overrides where input and

output types are different.

Args:

replacements (list[~apache_beam.pipeline.PTransformOverride]): a list of

:class:`~apache_beam.pipeline.PTransformOverride` objects.

"""

for override in replacements:

assert isinstance(override, PTransformOverride)

self._replace(override)

# Checking if the PTransforms have been successfully replaced. This will

# result in a failure if a PTransform that was replaced in a given override

# gets re-added in a subsequent override. This is not allowed and ordering

# of PTransformOverride objects in 'replacements' is important.

for override in replacements:

self._check_replacement(override)

def run(self, test_runner_api: Union[bool, str] = 'AUTO') -> 'PipelineResult':

"""Runs the pipeline. Returns whatever our runner returns after running."""

with scoped_pipeline_options(self._options):

return self._run_internal(test_runner_api)

def _run_internal(

self, test_runner_api: Union[bool, str] = 'AUTO') -> 'PipelineResult':

"""Internal implementation of run(), called within scoped options."""

# All pipeline options are finalized at this point.

# Call get_all_options to print warnings on invalid options.

self.options.get_all_options(

retain_unknown_options=True, display_warnings=True)

for error_handler in self._error_handlers:

error_handler.verify_closed()

# Records whether this pipeline contains any cross-language transforms.

self.contains_external_transforms = (

ExternalTransformFinder.contains_external_transforms(self))

try:

if test_runner_api == 'AUTO':

# Don't pay the cost of a round-trip if we're going to be going through

# the FnApi anyway...

is_fnapi_compatible = self.runner.is_fnapi_compatible() or (

# DirectRunner uses the Fn API for batch only

self.runner.__class__.__name__ == 'SwitchingDirectRunner' and

not self._options.view_as(StandardOptions).streaming)

# Multi-language pipelines that contain external pipeline segments may

# not be able to create a Python pipeline object graph. Hence following

# runner API check should be skipped for such pipelines.

# The InteractiveRunner relies on a constant pipeline reference, skip

# it.

test_runner_api = (

not is_fnapi_compatible and

not self.contains_external_transforms and

self.runner.__class__.__name__ != 'InteractiveRunner')

# When possible, invoke a round trip through the runner API.

if test_runner_api and self._verify_runner_api_compatible():

return Pipeline.from_runner_api(

self.to_runner_api(use_fake_coders=True),

self.runner,

self._options).run(False)

if (self._options.view_as(TypeOptions).runtime_type_check and

self._options.view_as(TypeOptions).performance_runtime_type_check):

raise RuntimeError(

'You cannot turn on runtime_type_check '

'and performance_runtime_type_check simultaneously. '

'Pick one or the other.')

if self._options.view_as(TypeOptions).runtime_type_check:

from apache_beam.typehints import typecheck

self.visit(typecheck.TypeCheckVisitor())

if self._options.view_as(TypeOptions).performance_runtime_type_check:

from apache_beam.typehints import typecheck

self.visit(typecheck.PerformanceTypeCheckVisitor())

if self._options.view_as(SetupOptions).save_main_session:

# If this option is chosen, verify we can pickle the main session early.

tmpdir = tempfile.mkdtemp()

try:

pickler.dump_session(os.path.join(tmpdir, 'main_session.pickle'))

finally:

shutil.rmtree(tmpdir)

return self.runner.run_pipeline(self, self._options)

finally:

if not is_in_ipython():

shutil.rmtree(self.local_tempdir, ignore_errors=True)

# else interactive beam handles the cleanup.

def __enter__(self) -> 'Pipeline':

self._extra_context = contextlib.ExitStack()

self._extra_context.enter_context(

subprocess_server.JavaJarServer.beam_services(

self._options.view_as(CrossLanguageOptions).beam_services))

self._extra_context.enter_context(

subprocess_server.SubprocessServer.cache_subprocesses())

return self

def __exit__(

self,

exc_type: Optional[Type[BaseException]],

exc_val: Optional[BaseException],

exc_tb: Optional['TracebackType']) -> None:

try:

if not exc_type:

self.result = self.run()

if not self._options.view_as(StandardOptions).no_wait_until_finish:

self.result.wait_until_finish()

else:

logging.info(

'Job execution continues without waiting for completion.'

' Use "wait_until_finish" in PipelineResult to block'

' until finished.')

finally:

self._extra_context.__exit__(exc_type, exc_val, exc_tb)

def visit(self, visitor: 'PipelineVisitor') -> None:

"""Visits depth-first every node of a pipeline's DAG.

Runner-internal implementation detail; no backwards-compatibility guarantees

Args:

visitor (~apache_beam.pipeline.PipelineVisitor):

:class:`~apache_beam.pipeline.PipelineVisitor` object whose callbacks

will be called for each node visited. See

:class:`~apache_beam.pipeline.PipelineVisitor` comments.

Raises:

TypeError: if node is specified and is not a

:class:`~apache_beam.pvalue.PValue`.

~apache_beam.error.PipelineError: if node is specified and does not

belong to this pipeline instance.

"""

visited: set[pvalue.PValue] = set()

self._root_transform().visit(visitor, self, visited)

def apply(

self,

transform: ptransform.PTransform,

pvalueish: Optional[pvalue.PValue] = None,

label: Optional[str] = None) -> pvalue.PValue:

"""Applies a custom transform using the pvalueish specified.

Args:

transform (~apache_beam.transforms.ptransform.PTransform): the

:class:`~apache_beam.transforms.ptransform.PTransform` to apply.

pvalueish (~apache_beam.pvalue.PCollection): the input for the

:class:`~apache_beam.transforms.ptransform.PTransform` (typically a

:class:`~apache_beam.pvalue.PCollection`).

label (str): label of the

:class:`~apache_beam.transforms.ptransform.PTransform`.

Raises:

TypeError: if the transform object extracted from the

argument list is not a

:class:`~apache_beam.transforms.ptransform.PTransform`.

RuntimeError: if the transform object was already applied to

this pipeline and needs to be cloned in order to apply again.

"""

with scoped_pipeline_options(self._options):

return self._apply_internal(transform, pvalueish, label)

def _apply_internal(

self,

transform: ptransform.PTransform,

pvalueish: Optional[pvalue.PValue] = None,

label: Optional[str] = None) -> pvalue.PValue:

"""Internal implementation of apply(), called within scoped options."""

if isinstance(transform, ptransform._NamedPTransform):

return self.apply(

transform.transform, pvalueish, label or transform.label)

if not label and isinstance(transform, ptransform._PTransformFnPTransform):

# This must be set before label is inspected.

transform.set_options(self._options)

if not isinstance(transform, ptransform.PTransform):

raise TypeError("Expected a PTransform object, got %s" % transform)

if label:

# Fix self.label as it is inspected by some PTransform operations

# (e.g. to produce error messages for type hint violations).

old_label, transform.label = transform.label, label

try:

return self.apply(transform, pvalueish)

finally:

transform.label = old_label

# Attempts to alter the label of the transform to be applied only when it's

# a top-level transform so that the cell number will not be prepended to

# every child transform in a composite.

if self._current_transform() is self._root_transform():

alter_label_if_ipython(transform, pvalueish)

full_label = '/'.join(

[self._current_transform().full_label, transform.label]).lstrip('/')

if full_label in self.applied_labels:

auto_unique_labels = self._options.view_as(

StandardOptions).auto_unique_labels

if auto_unique_labels:

# If auto_unique_labels is set, we will append a unique suffix to the

# label to make it unique.

logging.warning(

'Using --auto_unique_labels could cause data loss when '

'updating a pipeline or reloading the job state. '

'This is not recommended for streaming jobs.')

unique_label = self._generate_unique_label(transform)

return self.apply(transform, pvalueish, unique_label)

else:

raise RuntimeError(

'A transform with label "%s" already exists in the pipeline. '

'To apply a transform with a specified label, write '

'pvalue | "label" >> transform or use the option '

'"auto_unique_labels" to automatically generate unique '

'transform labels. Note "auto_unique_labels" '

'could cause data loss when updating a pipeline or '

'reloading the job state. This is not recommended for '

'streaming jobs.' % full_label)

self.applied_labels.add(full_label)

if pvalueish is None:

full_label = self._current_transform().full_label

raise TypeCheckError(

f'Transform "{full_label}" was applied to the output of '

f'an object of type None.')

pvalueish, inputs = transform._extract_input_pvalues(pvalueish)

try:

if not isinstance(inputs, dict):

inputs = {str(ix): input for (ix, input) in enumerate(inputs)}

except TypeError:

raise NotImplementedError(

'Unable to extract PValue inputs from %s; either %s does not accept '

'inputs of this format, or it does not properly override '

'_extract_input_pvalues' % (pvalueish, transform))

for t, leaf_input in inputs.items():

if not isinstance(leaf_input, pvalue.PValue) or not isinstance(t, str):

raise NotImplementedError(

'%s does not properly override _extract_input_pvalues, '

'returned %s from %s' % (transform, inputs, pvalueish))

current = AppliedPTransform(

self._current_transform(),

transform,

full_label,

inputs,

None,

annotations=self._current_annotations())

self._current_transform().add_part(current)

try:

self.transforms_stack.append(current)

type_options = self._options.view_as(TypeOptions)

if type_options.pipeline_type_check:

transform.type_check_inputs(pvalueish)

if isinstance(pvalueish, pvalue.PBegin) and isinstance(transform, ParDo):

full_label = self._current_transform().full_label

raise TypeCheckError(

f"Transform '{full_label}' expects a PCollection as input. "

"Got a PBegin/Pipeline instead.")

self._assert_not_applying_PDone(pvalueish, transform)

pvalueish_result = self.runner.apply(transform, pvalueish, self._options)

if type_options is not None and type_options.pipeline_type_check:

transform.type_check_outputs(pvalueish_result)

for tag, result in ptransform.get_named_nested_pvalues(pvalueish_result):

assert isinstance(result, (pvalue.PValue, pvalue.DoOutputsTuple))

# Make sure we set the producer only for a leaf node in the transform

# DAG. This way we preserve the last transform of a composite transform

# as being the real producer of the result.

if result.producer is None:

result.producer = current

# TODO(BEAM-1833): Pass full tuples dict.

self._infer_result_type(transform, tuple(inputs.values()), result)

assert isinstance(result.producer.inputs, tuple)

if isinstance(result, pvalue.DoOutputsTuple):

for tag, pc in list(result._pcolls.items()):

if tag not in current.outputs:

current.add_output(pc, tag)

continue

# If there is already a tag with the same name, increase a counter for

# the name. This can happen, for example, when a composite outputs a

# list of PCollections where all the tags are None.

base = tag

counter = 0

while tag in current.outputs:

counter += 1

tag = '%s_%d' % (base, counter)

current.add_output(result, tag)

if (type_options is not None and

type_options.type_check_strictness == 'ALL_REQUIRED' and

transform.get_type_hints().output_types is None):

ptransform_name = '%s(%s)' % (transform.__class__.__name__, full_label)

raise TypeCheckError(

'Pipeline type checking is enabled, however no '

'output type-hint was found for the '

'PTransform %s' % ptransform_name)

finally:

self.transforms_stack.pop()

return pvalueish_result

def _assert_not_applying_PDone(

self,

pvalueish: Optional[pvalue.PValue],

transform: ptransform.PTransform):

if isinstance(pvalueish, pvalue.PDone) and isinstance(transform, ParDo):

# If the input is a PDone, we cannot apply a ParDo transform.

full_label = self._current_transform().full_label

producer_label = pvalueish.producer.full_label

raise TypeCheckError(

f'Transform "{full_label}" was applied to the output of '

f'"{producer_label}" but "{producer_label.split("/")[-1]}" '

'produces no PCollections.')

def _generate_unique_label(self, transform: str) -> str:

"""

Given a transform, generate a unique label for it based on current label.

"""

unique_suffix = uuid.uuid4().hex[:6]

return '%s_%s' % (transform.label, unique_suffix)

def _infer_result_type(

self,

transform: ptransform.PTransform,

inputs: Sequence[Union[pvalue.PBegin, pvalue.PCollection]],

result_pcollection: Union[pvalue.PValue, pvalue.DoOutputsTuple]) -> None:

"""Infer and set the output element type for a PCollection.

This function determines the output types of transforms by combining:

1. Concrete input types from previous transforms

2. Type hints declared on the current transform

3. Type variable binding and substitution

TYPE VARIABLE BINDING

---------------------

Type variables (K, V, T, etc.) act as placeholders that get bound to

concrete types through pattern matching. This requires both an input

pattern and an output template:

Input Pattern (from .with_input_types()):

Defines where in the input to find each type variable

Example: Tuple[K, V] means "K is the first element, V is the second"

Output Template (from .with_output_types()):

Defines how to use the bound variables in the output

Example: Tuple[V, K] means "swap the positions"

CONCRETE TYPES VS TYPE VARIABLES

---------------------------------

The system handles these differently:

Concrete Types (e.g., str, int, Tuple[str, int]):

- Used as-is without any binding

- Do not fall back to Any

- Example: .with_output_types(Tuple[str, int]) → Tuple[str, int]

Type Variables (e.g., K, V, T):

- Must be bound through pattern matching

- Require .with_input_types() to provide the pattern

- Fall back to Any if not bound

- Example without pattern: Tuple[K, V] → Tuple[Any, Any]

- Example with pattern: Tuple[K, V] → Tuple[str, int]

BINDING ALGORITHM

-----------------

1. Match: Compare input pattern to concrete input

Pattern: Tuple[K, V]

Concrete: Tuple[str, int]

Result: {K: str, V: int} ← Bindings created

2. Substitute: Apply bindings to output template

Template: Tuple[V, K] ← Note: swapped!

Bindings: {K: str, V: int}

Result: Tuple[int, str] ← Swapped concrete types

Each transform operates in its own type inference scope. Type variables

declared in a parent composite transform do NOT automatically propagate

to child transforms.

Parent scope (composite):

@with_input_types(Tuple[K, V]) ← K, V defined here

class MyComposite(PTransform):

def expand(self, pcoll):

# Child scope - parent's K, V are NOT available

return pcoll | ChildTransform()

Type variables that remain unbound after inference fall back to Any:

EXAMPLES

--------

Example 1: Concrete types (no variables)

Input: Tuple[str, int]

Transform: .with_output_types(Tuple[str, int])

Output: Tuple[str, int] ← Used as-is

Example 2: Type variables with pattern (correct)

Input: Tuple[str, int]

Transform: .with_input_types(Tuple[K, V])

.with_output_types(Tuple[V, K])

Binding: {K: str, V: int}

Output: Tuple[int, str] ← Swapped!

Example 3: Type variables without pattern (falls back to Any)

Input: Tuple[str, int]

Transform: .with_output_types(Tuple[K, V]) ← No input pattern!

Binding: None (can't match)

Output: Tuple[Any, Any] ← Fallback

Example 4: Mixed concrete and variables

Input: Tuple[str, int]

Transform: .with_input_types(Tuple[str, V])

.with_output_types(Tuple[str, V])

Binding: {V: int} ← Only V needs binding

Output: Tuple[str, int] ← str passed through, V bound to int

Args:

transform: The PTransform being applied

inputs: Input PCollections (provides concrete types)

result_pcollection: Output PCollection to set type on

"""

# TODO(robertwb): Multi-input inference.

type_options = self._options.view_as(TypeOptions)

if type_options is None or not type_options.pipeline_type_check:

return

if (isinstance(result_pcollection, pvalue.PCollection) and

(not result_pcollection.element_type

# TODO(robertwb): Ideally we'd do intersection here.

or result_pcollection.element_type == typehints.Any)):

# {Single, multi}-input, single-output inference.

input_element_types_tuple = tuple(i.element_type for i in inputs)

input_element_type = (

input_element_types_tuple[0] if len(input_element_types_tuple) == 1

else typehints.Union[input_element_types_tuple])

type_hints = transform.get_type_hints()

if not result_pcollection.tag:

declared_output_type = type_hints.simple_output_type(transform.label)

else:

declared_output_type = type_hints.tagged_output_types().get(

result_pcollection.tag, typehints.Any)

if declared_output_type:

input_types = type_hints.input_types