Switch from namedtuple to frozen dataclass for TFT graph elements while adding more pytyping for easier maintenance.

PiperOrigin-RevId: 548621263

Author: zoyahav

tensorflow_transform/beam/analysis_graph_builder.py

@@ -14,9 +14,10 @@
 """Functions to create the implementation graph."""
 
 import collections
+import dataclasses
 import hashlib
-
-from typing import Dict, Mapping, Collection, Optional, Tuple
+from typing import Collection, Dict, Mapping, Optional, Tuple
+from typing import OrderedDict as OrderedDictType
 
 import tensorflow as tf
 from tensorflow_transform import analyzer_nodes
@@ -29,9 +30,6 @@
 from tensorflow_transform.beam import analyzer_cache
 from tensorflow_transform.beam import beam_nodes
 from tensorflow_transform.beam import combiner_packing_util
-# TODO(b/243513856): Switch to `collections.namedtuple` or `typing.NamedTuple`
-# once the Spark issue is resolved.
-from tfx_bsl.types import tfx_namedtuple
 
 
 # Used for debugging only. This will point to the most recent graph built.
@@ -109,11 +107,8 @@ def validate_value(self, value):
     assert isinstance(value, nodes.ValueNode)
 
 
-class _OptimizationView(
-    tfx_namedtuple.namedtuple('_OptimizationView', [
-        'prefer_fine_grained_view', 'flattened_view', 'fine_grained_view',
-        'hashed_path'
-    ])):
+@dataclasses.dataclass(frozen=True)
+class _OptimizationView:
   """A container for operation outputs during _OptimizeVisitor traversal.
 
   This is used in order to maintain both a flattened view, and a fine grained
@@ -123,29 +118,15 @@ class _OptimizationView(
   `fine_grained_view` should be used.  It should be set to true if the upstream
   view has cacheing operations that haven't been flattened yet.
   """
+  prefer_fine_grained_view: bool
+  flattened_view: nodes.ValueNode
+  fine_grained_view: Optional[OrderedDictType[str, nodes.ValueNode]]
+  hashed_path: Optional[bytes]
 
-  def __init__(self, prefer_fine_grained_view, flattened_view,
-               fine_grained_view, hashed_path):
-    if prefer_fine_grained_view and not fine_grained_view:
+  def __post_init__(self):
+    if self.prefer_fine_grained_view and not self.fine_grained_view:
       raise ValueError(
           'Cannot prefer fine_grained_view when one is not provided')
-    del hashed_path
-    self._validate_flattened_view(flattened_view)
-    self._validate_fine_grained_view(fine_grained_view)
-    super().__init__()
-
-  def _validate_flattened_view(self, view):
-    assert view is self.flattened_view
-    assert view is not None
-    assert isinstance(view, nodes.ValueNode), view
-
-  def _validate_fine_grained_view(self, view):
-    assert view is self.fine_grained_view
-    if view is None:
-      return
-    assert isinstance(view, collections.OrderedDict), view
-    for value in view.values():
-      assert isinstance(value, nodes.ValueNode), value
 
 
 class _OptimizeVisitor(nodes.Visitor):
@@ -287,8 +268,11 @@ def visit(self, operation_def, input_values):
       for view in input_values:
         disaggregated_input_values.extend(view.fine_grained_view.values())
 
-      # Checking that all cache has the same size.
-      assert len({len(value) for value in disaggregated_input_values}) == 1
+      # Each cache item should be a single ValueNode.
+      assert all(
+          isinstance(value, nodes.ValueNode)
+          for value in disaggregated_input_values
+      )
 
       next_inputs = nodes.apply_multi_output_operation(
           beam_nodes.Flatten,

tensorflow_transform/nodes.py

@@ -26,39 +26,10 @@
 
 import abc
 import collections
-from typing import Collection, Optional, Tuple
+import dataclasses
+from typing import Any, Collection, Dict, List, Optional, Tuple
 
 import pydot
-# TODO(b/243513856): Switch to `collections.namedtuple` or `typing.NamedTuple`
-# once the Spark issue is resolved.
-from tfx_bsl.types import tfx_namedtuple
-
-
-class ValueNode(
-    tfx_namedtuple.namedtuple('ValueNode',
-                              ['parent_operation', 'value_index'])):
-  """A placeholder that will ultimately be translated to a PCollection.
-
-  Attributes:
-    parent_operation: The `OperationNode` that produces this value.
-    value_index: The index of this value in the outputs of `parent_operation`.
-  """
-  __slots__ = ()
-
-  def __init__(self, parent_operation, value_index: int):
-    if not isinstance(parent_operation, OperationNode):
-      raise TypeError(
-          'parent_operation must be a OperationNode, got {} of type {}'.format(
-              parent_operation, type(parent_operation)))
-    num_outputs = parent_operation.operation_def.num_outputs
-    if not (0 <= value_index and value_index < num_outputs):
-      raise ValueError(
-          'value_index was {} but parent_operation had {} outputs'.format(
-              value_index, num_outputs))
-    super().__init__()
-
-  def __iter__(self):
-    raise ValueError('ValueNode is not iterable')
 
 
 class OperationDef(metaclass=abc.ABCMeta):
@@ -114,6 +85,28 @@ def cache_coder(self) -> Optional[object]:
     return None
 
 
+@dataclasses.dataclass(frozen=True)
+class ValueNode:
+  """A placeholder that will ultimately be translated to a PCollection.
+
+  Attributes:
+    parent_operation: The `OperationNode` that produces this value.
+    value_index: The index of this value in the outputs of `parent_operation`.
+  """
+
+  parent_operation: 'OperationNode'
+  value_index: int
+
+  def __post_init__(self):
+    num_outputs = self.parent_operation.operation_def.num_outputs
+    if not (0 <= self.value_index and self.value_index < num_outputs):
+      raise ValueError(
+          'value_index was {} but parent_operation had {} outputs'.format(
+              self.value_index, num_outputs
+          )
+      )
+
+
 class OperationNode:
   """A placeholder that will ultimately be translated to a PTransform.
 
@@ -229,8 +222,8 @@ def __init__(self, visitor: Visitor):
     Args:
       visitor: A `Visitor` object.
     """
-    self._cached_value_nodes_values = {}
-    self._stack = []
+    self._cached_value_nodes_values: Dict[ValueNode, Any] = {}
+    self._stack: List[OperationNode] = []
     self._visitor = visitor
 
   def visit_value_node(self, value_node: ValueNode):

tensorflow_transform/nodes_test.py

@@ -70,27 +70,6 @@ def testApplyOperationWithTupleOutput(self):
     self.assertEqual(op.inputs, (a, b))
     self.assertEqual(op.outputs, (b_copy, a_copy))
 
-  def testOperationNodeWithBadOperatonDef(self):
-    with self.assertRaisesRegexp(
-        TypeError, 'operation_def must be an OperationDef, got'):
-      nodes.OperationNode('not a operation_def', ())
-
-  def testOperationNodeWithBadInput(self):
-    a = nodes.apply_operation(_Constant, value='a', label='Constant[a]')
-    with self.assertRaisesRegexp(
-        TypeError, 'Inputs to Operation must be a ValueNode, got'):
-      nodes.OperationNode(_Concat(label='Concat'), (a, 'not a value_node'))
-
-  def testOperationNodeWithBadInputs(self):
-    with self.assertRaisesRegexp(
-        TypeError, 'inputs must be a tuple, got'):
-      nodes.OperationNode(_Concat(label='Concat'), 'not a tuple')
-
-  def testValueNodeWithBadParent(self):
-    with self.assertRaisesRegexp(
-        TypeError, 'parent_operation must be a OperationNode, got'):
-      nodes.ValueNode('not an operation node', 0)
-
   def testValueNodeWithNegativeValueIndex(self):
     a = nodes.apply_operation(_Constant, value='a', label='Constant[a]')
     with self.assertRaisesWithLiteralMatch(
@@ -183,7 +162,7 @@ def testTraverserOutputsNotATuple(self):
     mock_visitor = mock.MagicMock()
     mock_visitor.visit.side_effect = [42]
 
-    with self.assertRaisesRegexp(
+    with self.assertRaisesRegex(
         ValueError, r'expected visitor to return a tuple, got'):
       nodes.Traverser(mock_visitor).visit_value_node(a)
 
@@ -192,7 +171,7 @@ def testTraverserBadNumOutputs(self):
     mock_visitor = mock.MagicMock()
     mock_visitor.visit.side_effect = [('a', 'b')]
 
-    with self.assertRaisesRegexp(
+    with self.assertRaisesRegex(
         ValueError, 'has 1 outputs but visitor returned 2 values: '):
       nodes.Traverser(mock_visitor).visit_value_node(a)