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

"""Analyzes and modifies the pipeline that utilize the PCollection cache.

This module is experimental. No backwards-compatibility guarantees.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

import apache_beam as beam
from apache_beam.portability.api import beam_runner_api_pb2
from apache_beam.runners.interactive import cache_manager as cache


class PipelineAnalyzer(object):
  def __init__(self, cache_manager, pipeline_proto, underlying_runner,
               options=None, desired_cache_labels=None):
    """Constructor of PipelineAnanlyzer.

    Args:
      cache_manager: (CacheManager)
      pipeline_proto: (Pipeline proto)
      underlying_runner: (PipelineRunner)
      options: (PipelineOptions)
      desired_cache_labels: (Set[str]) a set of labels of the PCollection
        queried by the user.
    """
    self._cache_manager = cache_manager
    self._pipeline_proto = pipeline_proto
    self._desired_cache_labels = desired_cache_labels or []

    self._pipeline = beam.pipeline.Pipeline.from_runner_api(
        self._pipeline_proto,
        runner=underlying_runner,
        options=options)
    # context returned from to_runner_api is more informative than that returned
    # from from_runner_api.
    _, self._context = self._pipeline.to_runner_api(return_context=True)
    self._pipeline_info = PipelineInfo(self._pipeline_proto.components)

    # Result of the analysis that can be queried by the user.
    self._pipeline_proto_to_execute = None
    self._top_level_referenced_pcoll_ids = None
    self._top_level_required_transforms = None

    self._caches_used = set()
    self._read_cache_ids = set()
    self._write_cache_ids = set()

    # used for _insert_producing_transforms()
    self._analyzed_pcoll_ids = set()

    self._analyze_pipeline()

  def _analyze_pipeline(self):
    """Analyzes the pipeline and sets the variables that can be queried.

    This function construct Pipeline proto to execute by
      1. Start from target PCollections and recursively insert the producing
         PTransforms of those PCollections, where the producing PTransforms are
         either ReadCache or PTransforms in the original pipeline.
      2. Append WriteCache PTransforsm in the pipeline.

    After running this function, the following variables will be set:
      self._pipeline_proto_to_execute
      self._top_level_referenced_pcoll_ids
      self._top_level_required_transforms
      self._caches_used
      self._read_cache_ids
      self._write_cache_ids
    """
    # We filter PTransforms to be executed bottom-up from these PCollections.
    desired_pcollections = self._desired_pcollections(self._pipeline_info)

    required_transforms = collections.OrderedDict()
    top_level_required_transforms = collections.OrderedDict()

    for pcoll_id in desired_pcollections:
      # TODO(qinyeli): Collections consumed by no-output transforms.
      self._insert_producing_transforms(pcoll_id,
                                        required_transforms,
                                        top_level_required_transforms)

    top_level_referenced_pcoll_ids = self._referenced_pcoll_ids(
        top_level_required_transforms)

    for pcoll_id in self._pipeline_info.all_pcollections():
      if not pcoll_id in top_level_referenced_pcoll_ids:
        continue

      if (pcoll_id in desired_pcollections
          and not pcoll_id in self._caches_used):
        self._insert_caching_transforms(pcoll_id,
                                        required_transforms,
                                        top_level_required_transforms)

      if not self._cache_manager.exists(
          'sample', self._pipeline_info.cache_label(pcoll_id)):
        self._insert_caching_transforms(pcoll_id,
                                        required_transforms,
                                        top_level_required_transforms,
                                        sample=True)

    required_transforms['_root'] = beam_runner_api_pb2.PTransform(
        subtransforms=list(top_level_required_transforms.keys()))

    referenced_pcoll_ids = self._referenced_pcoll_ids(
        required_transforms)
    referenced_pcollections = {}
    for pcoll_id in referenced_pcoll_ids:
      obj = self._context.pcollections.get_by_id(pcoll_id)
      proto = self._context.pcollections.get_proto(obj)
      referenced_pcollections[pcoll_id] = proto

    pipeline_to_execute = beam_runner_api_pb2.Pipeline()
    pipeline_to_execute.root_transform_ids[:] = ['_root']
    set_proto_map(pipeline_to_execute.components.transforms,
                  required_transforms)
    set_proto_map(pipeline_to_execute.components.pcollections,
                  referenced_pcollections)
    set_proto_map(pipeline_to_execute.components.coders,
                  self._context.to_runner_api().coders)
    set_proto_map(pipeline_to_execute.components.windowing_strategies,
                  self._context.to_runner_api().windowing_strategies)

    self._pipeline_proto_to_execute = pipeline_to_execute
    self._top_level_referenced_pcoll_ids = top_level_referenced_pcoll_ids
    self._top_level_required_transforms = top_level_required_transforms

  # -------------------------------------------------------------------------- #
  # Getters
  # -------------------------------------------------------------------------- #

  def pipeline_info(self):
    """Return PipelineInfo of the original pipeline.
    """
    return self._pipeline_info

  def pipeline_proto_to_execute(self):
    """Returns Pipeline proto to be executed.
    """
    return self._pipeline_proto_to_execute

  def tl_referenced_pcoll_ids(self):
    """Returns a set of PCollection IDs referenced by top level PTransforms.
    """
    return self._top_level_referenced_pcoll_ids

  def tl_required_trans_ids(self):
    """Returns a set of required top level PTransform IDs.
    """
    return self._top_level_required_transforms.keys()

  def caches_used(self):
    """Returns a set of PCollection IDs to read from cache.
    """
    return self._caches_used

  def read_cache_ids(self):
    """Return a set of ReadCache PTransform IDs inserted.
    """
    return self._read_cache_ids

  def write_cache_ids(self):
    """Return a set of WriteCache PTransform IDs inserted.
    """
    return self._write_cache_ids

  # -------------------------------------------------------------------------- #
  # Helper methods for _analyze_pipeline()
  # -------------------------------------------------------------------------- #
<commit_msg>Fixing Py2-3 lint issue<commit_after>#
# 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.
#

"""Analyzes and modifies the pipeline that utilize the PCollection cache.

This module is experimental. No backwards-compatibility guarantees.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

import apache_beam as beam
from apache_beam.portability.api import beam_runner_api_pb2
from apache_beam.runners.interactive import cache_manager as cache


class PipelineAnalyzer(object):
  def __init__(self, cache_manager, pipeline_proto, underlying_runner,
               options=None, desired_cache_labels=None):
    """Constructor of PipelineAnanlyzer.

    Args:
      cache_manager: (CacheManager)
      pipeline_proto: (Pipeline proto)
      underlying_runner: (PipelineRunner)
      options: (PipelineOptions)
      desired_cache_labels: (Set[str]) a set of labels of the PCollection
        queried by the user.
    """
    self._cache_manager = cache_manager
    self._pipeline_proto = pipeline_proto
    self._desired_cache_labels = desired_cache_labels or []

    self._pipeline = beam.pipeline.Pipeline.from_runner_api(
        self._pipeline_proto,
        runner=underlying_runner,
        options=options)
    # context returned from to_runner_api is more informative than that returned
    # from from_runner_api.
    _, self._context = self._pipeline.to_runner_api(return_context=True)
    self._pipeline_info = PipelineInfo(self._pipeline_proto.components)

    # Result of the analysis that can be queried by the user.
    self._pipeline_proto_to_execute = None
    self._top_level_referenced_pcoll_ids = None
    self._top_level_required_transforms = None

    self._caches_used = set()
    self._read_cache_ids = set()
    self._write_cache_ids = set()

    # used for _insert_producing_transforms()
    self._analyzed_pcoll_ids = set()

    self._analyze_pipeline()

  def _analyze_pipeline(self):
    """Analyzes the pipeline and sets the variables that can be queried.

    This function construct Pipeline proto to execute by
      1. Start from target PCollections and recursively insert the producing
         PTransforms of those PCollections, where the producing PTransforms are
         either ReadCache or PTransforms in the original pipeline.
      2. Append WriteCache PTransforsm in the pipeline.

    After running this function, the following variables will be set:
      self._pipeline_proto_to_execute
      self._top_level_referenced_pcoll_ids
      self._top_level_required_transforms
      self._caches_used
      self._read_cache_ids
      self._write_cache_ids
    """
    # We filter PTransforms to be executed bottom-up from these PCollections.
    desired_pcollections = self._desired_pcollections(self._pipeline_info)

    required_transforms = collections.OrderedDict()
    top_level_required_transforms = collections.OrderedDict()

    for pcoll_id in desired_pcollections:
      # TODO(qinyeli): Collections consumed by no-output transforms.
      self._insert_producing_transforms(pcoll_id,
                                        required_transforms,
                                        top_level_required_transforms)

    top_level_referenced_pcoll_ids = self._referenced_pcoll_ids(
        top_level_required_transforms)

    for pcoll_id in self._pipeline_info.all_pcollections():
      if not pcoll_id in top_level_referenced_pcoll_ids:
        continue

      if (pcoll_id in desired_pcollections
          and not pcoll_id in self._caches_used):
        self._insert_caching_transforms(pcoll_id,
                                        required_transforms,
                                        top_level_required_transforms)

      if not self._cache_manager.exists(
          'sample', self._pipeline_info.cache_label(pcoll_id)):
        self._insert_caching_transforms(pcoll_id,
                                        required_transforms,
                                        top_level_required_transforms,
                                        sample=True)

    required_transforms['_root'] = beam_runner_api_pb2.PTransform(
        subtransforms=list(top_level_required_transforms))

    referenced_pcoll_ids = self._referenced_pcoll_ids(
        required_transforms)
    referenced_pcollections = {}
    for pcoll_id in referenced_pcoll_ids:
      obj = self._context.pcollections.get_by_id(pcoll_id)
      proto = self._context.pcollections.get_proto(obj)
      referenced_pcollections[pcoll_id] = proto

    pipeline_to_execute = beam_runner_api_pb2.Pipeline()
    pipeline_to_execute.root_transform_ids[:] = ['_root']
    set_proto_map(pipeline_to_execute.components.transforms,
                  required_transforms)
    set_proto_map(pipeline_to_execute.components.pcollections,
                  referenced_pcollections)
    set_proto_map(pipeline_to_execute.components.coders,
                  self._context.to_runner_api().coders)
    set_proto_map(pipeline_to_execute.components.windowing_strategies,
                  self._context.to_runner_api().windowing_strategies)

    self._pipeline_proto_to_execute = pipeline_to_execute
    self._top_level_referenced_pcoll_ids = top_level_referenced_pcoll_ids
    self._top_level_required_transforms = top_level_required_transforms

  # -------------------------------------------------------------------------- #
  # Getters
  # -------------------------------------------------------------------------- #

  def pipeline_info(self):
    """Return PipelineInfo of the original pipeline.
    """
    return self._pipeline_info

  def pipeline_proto_to_execute(self):
    """Returns Pipeline proto to be executed.
    """
    return self._pipeline_proto_to_execute

  def tl_referenced_pcoll_ids(self):
    """Returns a set of PCollection IDs referenced by top level PTransforms.
    """
    return self._top_level_referenced_pcoll_ids

  def tl_required_trans_ids(self):
    """Returns a set of required top level PTransform IDs.
    """
    return list(self._top_level_required_transforms)

  def caches_used(self):
    """Returns a set of PCollection IDs to read from cache.
    """
    return self._caches_used

  def read_cache_ids(self):
    """Return a set of ReadCache PTransform IDs inserted.
    """
    return self._read_cache_ids

  def write_cache_ids(self):
    """Return a set of WriteCache PTransform IDs inserted.
    """
    return self._write_cache_ids

  # -------------------------------------------------------------------------- #
  # Helper methods for _analyze_pipeline()
  # -------------------------------------------------------------------------- #

{'fraction_non_alphanumeric': 0.09835290941001755, 'fraction_numerical': 0.0013500742540839747, 'mean_word_length': 2.7996922287766095, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5325014', 'n_tokens_mistral': 4001, 'n_tokens_neox': 3952, 'n_words': 1023}