sdks/python/apache_beam/utils/windowed_value.py - beam - Git at Google

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

 """Core windowing data structures.

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

 # This module is carefully crafted to have optimal performance when
 # compiled while still being valid Python.  Care needs to be taken when
 # editing this file as WindowedValues are created for every element for
 # every step in a Beam pipeline.

 # cython: profile=True
 # cython: language_level=3

 # pytype: skip-file

 from typing import TYPE_CHECKING
 from typing import Any
 from typing import List
 from typing import Optional
 from typing import Tuple

 from apache_beam.utils.timestamp import MAX_TIMESTAMP
 from apache_beam.utils.timestamp import MIN_TIMESTAMP
 from apache_beam.utils.timestamp import Timestamp
 from apache_beam.utils.timestamp import TimestampTypes  # pylint: disable=unused-import

 if TYPE_CHECKING:
   from apache_beam.transforms.window import BoundedWindow


 class PaneInfoTiming(object):
   """The timing of a PaneInfo."""

   EARLY = 0
   ON_TIME = 1
   LATE = 2
   UNKNOWN = 3

   @classmethod
   def to_string(cls, value):
     return {
         cls.EARLY: 'EARLY',
         cls.ON_TIME: 'ON_TIME',
         cls.LATE: 'LATE',
         cls.UNKNOWN: 'UNKNOWN',
     }[value]

   @classmethod
   def from_string(cls, value):
     return {
         'EARLY': cls.EARLY,
         'ON_TIME': cls.ON_TIME,
         'LATE': cls.LATE,
         'UNKNOWN': cls.UNKNOWN
     }[value]


 class PaneInfo(object):
   """Describes the trigger firing information for a given WindowedValue.

   "Panes" represent individual firings on a single window. ``PaneInfo``s are
   passed downstream after trigger firings. They contain information about
   whether it's an early/on time/late firing, if it's the last or first firing
   from a window, and the index of the firing.
   """
   def __init__(self, is_first, is_last, timing, index, nonspeculative_index):
     self._is_first = is_first
     self._is_last = is_last
     self._timing = timing
     self._index = index
     self._nonspeculative_index = nonspeculative_index
     self._encoded_byte = self._get_encoded_byte()

   def _get_encoded_byte(self):
     byte = 0
     if self._is_first:
       byte |= 1
     if self._is_last:
       byte |= 2
     byte |= self._timing << 2
     return byte

   @staticmethod
   def from_encoded_byte(encoded_byte):
     assert encoded_byte in _BYTE_TO_PANE_INFO
     return _BYTE_TO_PANE_INFO[encoded_byte]

   # Because common PaneInfo objects are cached, it is important that the value
   # is immutable.  We therefore explicitly enforce this here with read-only
   # properties.

   @property
   def is_first(self):
     return self._is_first

   @property
   def is_last(self):
     return self._is_last

   @property
   def timing(self):
     return self._timing

   @property
   def index(self):
     # type: () -> int
     return self._index

   @property
   def nonspeculative_index(self):
     # type: () -> int
     return self._nonspeculative_index

   @property
   def encoded_byte(self):
     # type: () -> int
     return self._encoded_byte

   def __repr__(self):
     return (
         'PaneInfo(first: %r, last: %r, timing: %s, index: %d, '
         'nonspeculative_index: %d)') % (
             self.is_first,
             self.is_last,
             PaneInfoTiming.to_string(self.timing),
             self.index,
             self.nonspeculative_index)

   def __eq__(self, other):
     if self is other:
       return True
     return (
         self.is_first == other.is_first and self.is_last == other.is_last and
         self.timing == other.timing and self.index == other.index and
         self.nonspeculative_index == other.nonspeculative_index)

   def __hash__(self):
     return hash((
         self.is_first,
         self.is_last,
         self.timing,
         self.index,
         self.nonspeculative_index))

   def __reduce__(self):
     return PaneInfo, (self._is_first, self._is_last, self._timing, self._index,
                       self._nonspeculative_index)


 def _construct_well_known_pane_infos():
   # type: () -> List[PaneInfo]
   pane_infos = []
   for timing in (PaneInfoTiming.EARLY,
                  PaneInfoTiming.ON_TIME,
                  PaneInfoTiming.LATE,
                  PaneInfoTiming.UNKNOWN):
     nonspeculative_index = -1 if timing == PaneInfoTiming.EARLY else 0
     pane_infos.append(PaneInfo(True, True, timing, 0, nonspeculative_index))
     pane_infos.append(PaneInfo(True, False, timing, 0, nonspeculative_index))
     pane_infos.append(PaneInfo(False, True, timing, -1, nonspeculative_index))
     pane_infos.append(PaneInfo(False, False, timing, -1, nonspeculative_index))
   result = [None] * (
       max(p.encoded_byte for p in pane_infos) + 1
   )  # type: List[PaneInfo]  # type: ignore[list-item]
   for pane_info in pane_infos:
     result[pane_info.encoded_byte] = pane_info
   return result


 # Cache of well-known PaneInfo objects.
 _BYTE_TO_PANE_INFO = _construct_well_known_pane_infos()

 # Default PaneInfo descriptor for when a value is not the output of triggering.
 PANE_INFO_UNKNOWN = _BYTE_TO_PANE_INFO[0xF]


 class WindowedValue(object):
   """A windowed value having a value, a timestamp and set of windows.

   Attributes:
     value: The underlying value of a windowed value.
     timestamp: Timestamp associated with the value as seconds since Unix epoch.
     windows: A set (iterable) of window objects for the value. The window
       object are descendants of the BoundedWindow class.
     pane_info: A PaneInfo descriptor describing the triggering information for
       the pane that contained this value.  If None, will be set to
       PANE_INFO_UNKNOWN.
   """

   def __init__(self,
                value,
                timestamp,  # type: TimestampTypes
                windows,  # type: Tuple[BoundedWindow, ...]
                pane_info=PANE_INFO_UNKNOWN  # type: PaneInfo
               ):
     # type: (...) -> None
     # For performance reasons, only timestamp_micros is stored by default
     # (as a C int). The Timestamp object is created on demand below.
     self.value = value
     if isinstance(timestamp, int):
       self.timestamp_micros = timestamp * 1000000
       if TYPE_CHECKING:
         self.timestamp_object = None  # type: Optional[Timestamp]
     else:
       self.timestamp_object = (
           timestamp
           if isinstance(timestamp, Timestamp) else Timestamp.of(timestamp))
       self.timestamp_micros = self.timestamp_object.micros
     self.windows = windows
     self.pane_info = pane_info

   @property
   def timestamp(self):
     # type: () -> Timestamp
     if self.timestamp_object is None:
       self.timestamp_object = Timestamp(0, self.timestamp_micros)
     return self.timestamp_object

   def __repr__(self):
     return '(%s, %s, %s, %s)' % (
         repr(self.value),
         'MIN_TIMESTAMP' if self.timestamp == MIN_TIMESTAMP else 'MAX_TIMESTAMP'
         if self.timestamp == MAX_TIMESTAMP else float(self.timestamp),
         self.windows,
         self.pane_info)

   def __eq__(self, other):
     return (
         type(self) == type(other) and
         self.timestamp_micros == other.timestamp_micros and
         self.value == other.value and self.windows == other.windows and
         self.pane_info == other.pane_info)

   def __hash__(self):
     return ((hash(self.value) & 0xFFFFFFFFFFFFFFF) + 3 *
             (self.timestamp_micros & 0xFFFFFFFFFFFFFF) + 7 *
             (hash(self.windows) & 0xFFFFFFFFFFFFF) + 11 *
             (hash(self.pane_info) & 0xFFFFFFFFFFFFF))

   def with_value(self, new_value):
     # type: (Any) -> WindowedValue

     """Creates a new WindowedValue with the same timestamps and windows as this.

     This is the fasted way to create a new WindowedValue.
     """
     return create(
         new_value, self.timestamp_micros, self.windows, self.pane_info)

   def __reduce__(self):
     return WindowedValue, (
         self.value, self.timestamp, self.windows, self.pane_info)


 # TODO(robertwb): Move this to a static method.
 def create(value, timestamp_micros, windows, pane_info=PANE_INFO_UNKNOWN):
   wv = WindowedValue.__new__(WindowedValue)
   wv.value = value
   wv.timestamp_micros = timestamp_micros
   wv.windows = windows
   wv.pane_info = pane_info
   return wv


 try:
   WindowedValue.timestamp_object = None
 except TypeError:
   # When we're compiled, we can't dynamically add attributes to
   # the cdef class, but in this case it's OK as it's already present
   # on each instance.
   pass


 class _IntervalWindowBase(object):
   """Optimized form of IntervalWindow storing only microseconds for endpoints.
   """
   def __init__(self, start, end):
     # type: (TimestampTypes, TimestampTypes) -> None
     if start is not None:
       self._start_object = Timestamp.of(start)  # type: Optional[Timestamp]
       try:
         self._start_micros = self._start_object.micros
       except OverflowError:
         self._start_micros = (
             MIN_TIMESTAMP.micros
             if self._start_object.micros < 0 else MAX_TIMESTAMP.micros)
     else:
       # Micros must be populated elsewhere.
       self._start_object = None

     if end is not None:
       self._end_object = Timestamp.of(end)  # type: Optional[Timestamp]
       try:
         self._end_micros = self._end_object.micros
       except OverflowError:
         self._end_micros = (
             MIN_TIMESTAMP.micros
             if self._end_object.micros < 0 else MAX_TIMESTAMP.micros)
     else:
       # Micros must be populated elsewhere.
       self._end_object = None

   @property
   def start(self):
     # type: () -> Timestamp
     if self._start_object is None:
       self._start_object = Timestamp(0, self._start_micros)
     return self._start_object

   @property
   def end(self):
     # type: () -> Timestamp
     if self._end_object is None:
       self._end_object = Timestamp(0, self._end_micros)
     return self._end_object

   def __hash__(self):
     return hash((self._start_micros, self._end_micros))

   def __eq__(self, other):
     return (
         type(self) == type(other) and
         self._start_micros == other._start_micros and
         self._end_micros == other._end_micros)

   def __repr__(self):
     return '[%s, %s)' % (float(self.start), float(self.end))
	#
	# 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.
	#

	"""Core windowing data structures.

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

	# This module is carefully crafted to have optimal performance when
	# compiled while still being valid Python. Care needs to be taken when
	# editing this file as WindowedValues are created for every element for
	# every step in a Beam pipeline.

	# cython: profile=True
	# cython: language_level=3

	# pytype: skip-file

	from typing import TYPE_CHECKING
	from typing import Any
	from typing import List
	from typing import Optional
	from typing import Tuple

	from apache_beam.utils.timestamp import MAX_TIMESTAMP
	from apache_beam.utils.timestamp import MIN_TIMESTAMP
	from apache_beam.utils.timestamp import Timestamp
	from apache_beam.utils.timestamp import TimestampTypes # pylint: disable=unused-import

	if TYPE_CHECKING:
	from apache_beam.transforms.window import BoundedWindow


	class PaneInfoTiming(object):
	"""The timing of a PaneInfo."""

	EARLY = 0
	ON_TIME = 1
	LATE = 2
	UNKNOWN = 3

	@classmethod
	def to_string(cls, value):
	return {
	cls.EARLY: 'EARLY',
	cls.ON_TIME: 'ON_TIME',
	cls.LATE: 'LATE',
	cls.UNKNOWN: 'UNKNOWN',
	}[value]

	@classmethod
	def from_string(cls, value):
	return {
	'EARLY': cls.EARLY,
	'ON_TIME': cls.ON_TIME,
	'LATE': cls.LATE,
	'UNKNOWN': cls.UNKNOWN
	}[value]


	class PaneInfo(object):
	"""Describes the trigger firing information for a given WindowedValue.

	"Panes" represent individual firings on a single window. ``PaneInfo``s are
	passed downstream after trigger firings. They contain information about
	whether it's an early/on time/late firing, if it's the last or first firing
	from a window, and the index of the firing.
	"""
	def __init__(self, is_first, is_last, timing, index, nonspeculative_index):
	self._is_first = is_first
	self._is_last = is_last
	self._timing = timing
	self._index = index
	self._nonspeculative_index = nonspeculative_index
	self._encoded_byte = self._get_encoded_byte()

	def _get_encoded_byte(self):
	byte = 0
	if self._is_first:
	byte \|= 1
	if self._is_last:
	byte \|= 2
	byte \|= self._timing << 2
	return byte

	@staticmethod
	def from_encoded_byte(encoded_byte):
	assert encoded_byte in _BYTE_TO_PANE_INFO
	return _BYTE_TO_PANE_INFO[encoded_byte]

	# Because common PaneInfo objects are cached, it is important that the value
	# is immutable. We therefore explicitly enforce this here with read-only
	# properties.

	@property
	def is_first(self):
	return self._is_first

	@property
	def is_last(self):
	return self._is_last

	@property
	def timing(self):
	return self._timing

	@property
	def index(self):
	# type: () -> int
	return self._index

	@property
	def nonspeculative_index(self):
	# type: () -> int
	return self._nonspeculative_index

	@property
	def encoded_byte(self):
	# type: () -> int
	return self._encoded_byte

	def __repr__(self):
	return (
	'PaneInfo(first: %r, last: %r, timing: %s, index: %d, '
	'nonspeculative_index: %d)') % (
	self.is_first,
	self.is_last,
	PaneInfoTiming.to_string(self.timing),
	self.index,
	self.nonspeculative_index)

	def __eq__(self, other):
	if self is other:
	return True
	return (
	self.is_first == other.is_first and self.is_last == other.is_last and
	self.timing == other.timing and self.index == other.index and
	self.nonspeculative_index == other.nonspeculative_index)

	def __hash__(self):
	return hash((
	self.is_first,
	self.is_last,
	self.timing,
	self.index,
	self.nonspeculative_index))

	def __reduce__(self):
	return PaneInfo, (self._is_first, self._is_last, self._timing, self._index,
	self._nonspeculative_index)


	def _construct_well_known_pane_infos():
	# type: () -> List[PaneInfo]
	pane_infos = []
	for timing in (PaneInfoTiming.EARLY,
	PaneInfoTiming.ON_TIME,
	PaneInfoTiming.LATE,
	PaneInfoTiming.UNKNOWN):
	nonspeculative_index = -1 if timing == PaneInfoTiming.EARLY else 0
	pane_infos.append(PaneInfo(True, True, timing, 0, nonspeculative_index))
	pane_infos.append(PaneInfo(True, False, timing, 0, nonspeculative_index))
	pane_infos.append(PaneInfo(False, True, timing, -1, nonspeculative_index))
	pane_infos.append(PaneInfo(False, False, timing, -1, nonspeculative_index))
	result = [None] * (
	max(p.encoded_byte for p in pane_infos) + 1
	) # type: List[PaneInfo] # type: ignore[list-item]
	for pane_info in pane_infos:
	result[pane_info.encoded_byte] = pane_info
	return result


	# Cache of well-known PaneInfo objects.
	_BYTE_TO_PANE_INFO = _construct_well_known_pane_infos()

	# Default PaneInfo descriptor for when a value is not the output of triggering.
	PANE_INFO_UNKNOWN = _BYTE_TO_PANE_INFO[0xF]


	class WindowedValue(object):
	"""A windowed value having a value, a timestamp and set of windows.

	Attributes:
	value: The underlying value of a windowed value.
	timestamp: Timestamp associated with the value as seconds since Unix epoch.
	windows: A set (iterable) of window objects for the value. The window
	object are descendants of the BoundedWindow class.
	pane_info: A PaneInfo descriptor describing the triggering information for
	the pane that contained this value. If None, will be set to
	PANE_INFO_UNKNOWN.
	"""

	def __init__(self,
	value,
	timestamp, # type: TimestampTypes
	windows, # type: Tuple[BoundedWindow, ...]
	pane_info=PANE_INFO_UNKNOWN # type: PaneInfo
	):
	# type: (...) -> None
	# For performance reasons, only timestamp_micros is stored by default
	# (as a C int). The Timestamp object is created on demand below.
	self.value = value
	if isinstance(timestamp, int):
	self.timestamp_micros = timestamp * 1000000
	if TYPE_CHECKING:
	self.timestamp_object = None # type: Optional[Timestamp]
	else:
	self.timestamp_object = (
	timestamp
	if isinstance(timestamp, Timestamp) else Timestamp.of(timestamp))
	self.timestamp_micros = self.timestamp_object.micros
	self.windows = windows
	self.pane_info = pane_info

	@property
	def timestamp(self):
	# type: () -> Timestamp
	if self.timestamp_object is None:
	self.timestamp_object = Timestamp(0, self.timestamp_micros)
	return self.timestamp_object

	def __repr__(self):
	return '(%s, %s, %s, %s)' % (
	repr(self.value),
	'MIN_TIMESTAMP' if self.timestamp == MIN_TIMESTAMP else 'MAX_TIMESTAMP'
	if self.timestamp == MAX_TIMESTAMP else float(self.timestamp),
	self.windows,
	self.pane_info)

	def __eq__(self, other):
	return (
	type(self) == type(other) and
	self.timestamp_micros == other.timestamp_micros and
	self.value == other.value and self.windows == other.windows and
	self.pane_info == other.pane_info)

	def __hash__(self):
	return ((hash(self.value) & 0xFFFFFFFFFFFFFFF) + 3 *
	(self.timestamp_micros & 0xFFFFFFFFFFFFFF) + 7 *
	(hash(self.windows) & 0xFFFFFFFFFFFFF) + 11 *
	(hash(self.pane_info) & 0xFFFFFFFFFFFFF))

	def with_value(self, new_value):
	# type: (Any) -> WindowedValue

	"""Creates a new WindowedValue with the same timestamps and windows as this.

	This is the fasted way to create a new WindowedValue.
	"""
	return create(
	new_value, self.timestamp_micros, self.windows, self.pane_info)

	def __reduce__(self):
	return WindowedValue, (
	self.value, self.timestamp, self.windows, self.pane_info)


	# TODO(robertwb): Move this to a static method.
	def create(value, timestamp_micros, windows, pane_info=PANE_INFO_UNKNOWN):
	wv = WindowedValue.__new__(WindowedValue)
	wv.value = value
	wv.timestamp_micros = timestamp_micros
	wv.windows = windows
	wv.pane_info = pane_info
	return wv


	try:
	WindowedValue.timestamp_object = None
	except TypeError:
	# When we're compiled, we can't dynamically add attributes to
	# the cdef class, but in this case it's OK as it's already present
	# on each instance.
	pass


	class _IntervalWindowBase(object):
	"""Optimized form of IntervalWindow storing only microseconds for endpoints.
	"""
	def __init__(self, start, end):
	# type: (TimestampTypes, TimestampTypes) -> None
	if start is not None:
	self._start_object = Timestamp.of(start) # type: Optional[Timestamp]
	try:
	self._start_micros = self._start_object.micros
	except OverflowError:
	self._start_micros = (
	MIN_TIMESTAMP.micros
	if self._start_object.micros < 0 else MAX_TIMESTAMP.micros)
	else:
	# Micros must be populated elsewhere.
	self._start_object = None

	if end is not None:
	self._end_object = Timestamp.of(end) # type: Optional[Timestamp]
	try:
	self._end_micros = self._end_object.micros
	except OverflowError:
	self._end_micros = (
	MIN_TIMESTAMP.micros
	if self._end_object.micros < 0 else MAX_TIMESTAMP.micros)
	else:
	# Micros must be populated elsewhere.
	self._end_object = None

	@property
	def start(self):
	# type: () -> Timestamp
	if self._start_object is None:
	self._start_object = Timestamp(0, self._start_micros)
	return self._start_object

	@property
	def end(self):
	# type: () -> Timestamp
	if self._end_object is None:
	self._end_object = Timestamp(0, self._end_micros)
	return self._end_object

	def __hash__(self):
	return hash((self._start_micros, self._end_micros))

	def __eq__(self, other):
	return (
	type(self) == type(other) and
	self._start_micros == other._start_micros and
	self._end_micros == other._end_micros)

	def __repr__(self):
	return '[%s, %s)' % (float(self.start), float(self.end))