github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/runners/worker/sideinputs.py

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

"""Utilities for handling side inputs."""

# pytype: skip-file

import logging
import queue
import threading
import traceback
from collections import abc

from apache_beam.coders import observable
from apache_beam.io import iobase
from apache_beam.runners.worker import opcounters
from apache_beam.transforms import window
from apache_beam.utils.sentinel import Sentinel

# Maximum number of reader threads for reading side input sources, per side
# input.
MAX_SOURCE_READER_THREADS = 15

# Number of slots for elements in side input element queue.  Note that this
# value is intentionally smaller than MAX_SOURCE_READER_THREADS so as to reduce
# memory pressure of holding potentially-large elements in memory.  Note that
# the number of pending elements in memory is equal to the sum of
# MAX_SOURCE_READER_THREADS and ELEMENT_QUEUE_SIZE.
ELEMENT_QUEUE_SIZE = 10

# Special element value sentinel for signaling reader state.
READER_THREAD_IS_DONE_SENTINEL = Sentinel.sentinel

# Used to efficiently window the values of non-windowed side inputs.
_globally_windowed = window.GlobalWindows.windowed_value(None).with_value

_LOGGER = logging.getLogger(__name__)


class PrefetchingSourceSetIterable(object):
  """Value iterator that reads concurrently from a set of sources."""
  def __init__(
      self,
      sources,
      max_reader_threads=MAX_SOURCE_READER_THREADS,
      read_counter=None,
      element_counter=None):
    self.sources = sources
    self.num_reader_threads = min(max_reader_threads, len(self.sources))

    # Queue for sources that are to be read.
    self.sources_queue = queue.Queue()
    for source in sources:
      self.sources_queue.put(source)
    # Queue for elements that have been read.
    self.element_queue = queue.Queue(ELEMENT_QUEUE_SIZE)
    # Queue for exceptions encountered in reader threads; to be rethrown.
    self.reader_exceptions = queue.Queue()
    # Whether we have already iterated; this iterable can only be used once.
    self.already_iterated = False
    # Whether an error was encountered in any source reader.
    self.has_errored = False

    self.read_counter = read_counter or opcounters.NoOpTransformIOCounter()
    self.element_counter = element_counter
    self.reader_threads = []
    self._start_reader_threads()

  def add_byte_counter(self, reader):
    """Adds byte counter observer to a side input reader.

    Args:
      reader: A reader that should inherit from ObservableMixin to have
        bytes tracked.
    """
    def update_bytes_read(record_size, is_record_size=False, **kwargs):
      # Let the reader report block size.
      if is_record_size:
        self.read_counter.add_bytes_read(record_size)

    if isinstance(reader, observable.ObservableMixin):
      reader.register_observer(update_bytes_read)

  def _start_reader_threads(self):
    for _ in range(0, self.num_reader_threads):
      t = threading.Thread(target=self._reader_thread)
      t.daemon = True
      t.start()
      self.reader_threads.append(t)

  def _reader_thread(self):
    # pylint: disable=too-many-nested-blocks
    try:
      while True:
        try:
          source = self.sources_queue.get_nowait()
          if isinstance(source, iobase.BoundedSource):
            for value in source.read(source.get_range_tracker(None, None)):
              if self.has_errored:
                # If any reader has errored, just return.
                return
              if isinstance(value, window.WindowedValue):
                self.element_queue.put(value)
              else:
                self.element_queue.put(_globally_windowed(value))
          else:
            # Native dataflow source.
            with source.reader() as reader:
              # The tracking of time spend reading and bytes read from side
              # inputs is kept behind an experiment flag to test performance
              # impact.
              self.add_byte_counter(reader)
              returns_windowed_values = reader.returns_windowed_values
              for value in reader:
                if self.has_errored:
                  # If any reader has errored, just return.
                  return
                if returns_windowed_values:
                  self.element_queue.put(value)
                else:
                  self.element_queue.put(_globally_windowed(value))
        except queue.Empty:
          return
    except Exception as e:  # pylint: disable=broad-except
      _LOGGER.error(
          'Encountered exception in PrefetchingSourceSetIterable '
          'reader thread: %s',
          traceback.format_exc())
      self.reader_exceptions.put(e)
      self.has_errored = True
    finally:
      self.element_queue.put(READER_THREAD_IS_DONE_SENTINEL)

  def __iter__(self):
    # pylint: disable=too-many-nested-blocks
    if self.already_iterated:
      raise RuntimeError(
          'Can only iterate once over PrefetchingSourceSetIterable instance.')
    self.already_iterated = True

    # The invariants during execution are:
    # 1) A worker thread always posts the sentinel as the last thing it does
    #    before exiting.
    # 2) We always wait for all sentinels and then join all threads.
    num_readers_finished = 0
    try:
      while True:
        try:
          with self.read_counter:
            element = self.element_queue.get()
          if element is READER_THREAD_IS_DONE_SENTINEL:
            num_readers_finished += 1
            if num_readers_finished == self.num_reader_threads:
              return
          else:
            if self.element_counter:
              self.element_counter.update_from(element)
              yield element
              self.element_counter.update_collect()
            else:
              yield element
        finally:
          if self.has_errored:
            raise self.reader_exceptions.get()
    except GeneratorExit:
      self.has_errored = True
      raise
    finally:
      while num_readers_finished < self.num_reader_threads:
        element = self.element_queue.get()
        if element is READER_THREAD_IS_DONE_SENTINEL:
          num_readers_finished += 1
      for t in self.reader_threads:
        t.join()


def get_iterator_fn_for_sources(
    sources,
    max_reader_threads=MAX_SOURCE_READER_THREADS,
    read_counter=None,
    element_counter=None):
  """Returns callable that returns iterator over elements for given sources."""
  def _inner():
    return iter(
        PrefetchingSourceSetIterable(
            sources,
            max_reader_threads=max_reader_threads,
            read_counter=read_counter,
            element_counter=element_counter))

  return _inner


class EmulatedIterable(abc.Iterable):
  """Emulates an iterable for a side input."""
  def __init__(self, iterator_fn):
    self.iterator_fn = iterator_fn

  def __iter__(self):
    return self.iterator_fn()