n/a

PiperOrigin-RevId: 522298075

Author: zoyahav

examples/simple_example_test.py

@@ -14,8 +14,8 @@
 """Tests for simple_example."""
 
 import tensorflow as tf
-from tensorflow_transform.beam import tft_unit
 import simple_example
+from tensorflow_transform.beam import tft_unit
 
 
 _EXPECTED_TRANSFORMED_OUTPUT = [

examples/simple_sequence_example_test.py

@@ -14,8 +14,8 @@
 """Tests for simple_example."""
 
 import tensorflow as tf
-from tensorflow_transform.beam import tft_unit
 import simple_sequence_example
+from tensorflow_transform.beam import tft_unit
 
 _EXPECTED_TRANSFORMED_OUTPUT = [{
     'transformed_seq_int_feature$ragged_values': [

tensorflow_transform/beam/analysis_graph_builder_test.py

@@ -23,7 +23,7 @@
 from tensorflow_transform import tf2_utils
 from tensorflow_transform.beam import analysis_graph_builder
 from tensorflow_transform.beam import analyzer_cache
-from tensorflow_transform import test_case
+from tensorflow_transform.beam import tft_unit
 # TODO(b/243513856): Switch to `collections.namedtuple` or `typing.NamedTuple`
 # once the Spark issue is resolved.
 from tfx_bsl.types import tfx_namedtuple
@@ -396,17 +396,21 @@ def __new__(cls):
 ]
 
 
-class AnalysisGraphBuilderTest(test_case.TransformTestCase):
+class AnalysisGraphBuilderTest(tft_unit.TransformTestCase):
 
-  @test_case.named_parameters(
-      *test_case.cross_named_parameters(_ANALYZE_TEST_CASES, [
-          dict(testcase_name='tf_compat_v1', use_tf_compat_v1=True),
-          dict(testcase_name='tf2', use_tf_compat_v1=False)
-      ]))
+  @tft_unit.named_parameters(
+      *tft_unit.cross_named_parameters(
+          _ANALYZE_TEST_CASES,
+          [
+              dict(testcase_name='tf_compat_v1', use_tf_compat_v1=True),
+              dict(testcase_name='tf2', use_tf_compat_v1=False),
+          ],
+      )
+  )
   def test_build(self, feature_spec, preprocessing_fn, expected_dot_graph_str,
                  expected_dot_graph_str_tf2, use_tf_compat_v1):
     if not use_tf_compat_v1:
-      test_case.skip_if_not_tf2('Tensorflow 2.x required')
+      tft_unit.skip_if_not_tf2('Tensorflow 2.x required')
     specs = (
         feature_spec if use_tf_compat_v1 else
         impl_helper.get_type_specs_from_feature_specs(feature_spec))
@@ -430,48 +434,54 @@ def test_build(self, feature_spec, preprocessing_fn, expected_dot_graph_str,
         second=(expected_dot_graph_str
                 if use_tf_compat_v1 else expected_dot_graph_str_tf2))
 
-  @test_case.named_parameters(*test_case.cross_named_parameters(
-      [
-          dict(
-              testcase_name='one_dataset_cached_single_phase',
-              preprocessing_fn=_preprocessing_fn_with_one_analyzer,
-              full_dataset_keys=['a', 'b'],
-              cached_dataset_keys=['a'],
-              expected_dataset_keys=['b'],
-          ),
-          dict(
-              testcase_name='all_datasets_cached_single_phase',
-              preprocessing_fn=_preprocessing_fn_with_one_analyzer,
-              full_dataset_keys=['a', 'b'],
-              cached_dataset_keys=['a', 'b'],
-              expected_dataset_keys=[],
-          ),
-          dict(
-              testcase_name='mixed_single_phase',
-              preprocessing_fn=lambda d: dict(  # pylint: disable=g-long-lambda
-                  list(_preprocessing_fn_with_chained_ptransforms(d).items()) +
-                  list(_preprocessing_fn_with_one_analyzer(d).items())),
-              full_dataset_keys=['a', 'b'],
-              cached_dataset_keys=['a', 'b'],
-              expected_dataset_keys=['a', 'b'],
-          ),
-          dict(
-              testcase_name='multi_phase',
-              preprocessing_fn=_preprocessing_fn_with_two_phases,
-              full_dataset_keys=['a', 'b'],
-              cached_dataset_keys=['a', 'b'],
-              expected_dataset_keys=['a', 'b'],
-          )
-      ],
-      [
-          dict(testcase_name='tf_compat_v1', use_tf_compat_v1=True),
-          dict(testcase_name='tf2', use_tf_compat_v1=False)
-      ]))
+  @tft_unit.named_parameters(
+      *tft_unit.cross_named_parameters(
+          [
+              dict(
+                  testcase_name='one_dataset_cached_single_phase',
+                  preprocessing_fn=_preprocessing_fn_with_one_analyzer,
+                  full_dataset_keys=['a', 'b'],
+                  cached_dataset_keys=['a'],
+                  expected_dataset_keys=['b'],
+              ),
+              dict(
+                  testcase_name='all_datasets_cached_single_phase',
+                  preprocessing_fn=_preprocessing_fn_with_one_analyzer,
+                  full_dataset_keys=['a', 'b'],
+                  cached_dataset_keys=['a', 'b'],
+                  expected_dataset_keys=[],
+              ),
+              dict(
+                  testcase_name='mixed_single_phase',
+                  preprocessing_fn=lambda d: dict(  # pylint: disable=g-long-lambda
+                      list(
+                          _preprocessing_fn_with_chained_ptransforms(d).items()
+                      )
+                      + list(_preprocessing_fn_with_one_analyzer(d).items())
+                  ),
+                  full_dataset_keys=['a', 'b'],
+                  cached_dataset_keys=['a', 'b'],
+                  expected_dataset_keys=['a', 'b'],
+              ),
+              dict(
+                  testcase_name='multi_phase',
+                  preprocessing_fn=_preprocessing_fn_with_two_phases,
+                  full_dataset_keys=['a', 'b'],
+                  cached_dataset_keys=['a', 'b'],
+                  expected_dataset_keys=['a', 'b'],
+              ),
+          ],
+          [
+              dict(testcase_name='tf_compat_v1', use_tf_compat_v1=True),
+              dict(testcase_name='tf2', use_tf_compat_v1=False),
+          ],
+      )
+  )
   def test_get_analysis_dataset_keys(self, preprocessing_fn, full_dataset_keys,
                                      cached_dataset_keys, expected_dataset_keys,
                                      use_tf_compat_v1):
     if not use_tf_compat_v1:
-      test_case.skip_if_not_tf2('Tensorflow 2.x required')
+      tft_unit.skip_if_not_tf2('Tensorflow 2.x required')
     full_dataset_keys = list(
         map(analyzer_cache.DatasetKey, full_dataset_keys))
     cached_dataset_keys = map(analyzer_cache.DatasetKey, cached_dataset_keys)
@@ -499,18 +509,16 @@ def test_get_analysis_dataset_keys(self, preprocessing_fn, full_dataset_keys,
               full_dataset_keys,
               input_cache,
               force_tf_compat_v1=use_tf_compat_v1))
-
-    dot_string = nodes.get_dot_graph([analysis_graph_builder._ANALYSIS_GRAPH
-                                     ]).to_string()
-    self.WriteRenderedDotFile(dot_string)
+    self.DebugPublishLatestsRenderedTFTGraph()
     self.assertCountEqual(expected_dataset_keys, dataset_keys)
 
-  @test_case.named_parameters(
+  @tft_unit.named_parameters(
       dict(testcase_name='tf_compat_v1', use_tf_compat_v1=True),
-      dict(testcase_name='tf2', use_tf_compat_v1=False))
+      dict(testcase_name='tf2', use_tf_compat_v1=False),
+  )
   def test_get_analysis_cache_entry_keys(self, use_tf_compat_v1):
     if not use_tf_compat_v1:
-      test_case.skip_if_not_tf2('Tensorflow 2.x required')
+      tft_unit.skip_if_not_tf2('Tensorflow 2.x required')
     full_dataset_keys = map(analyzer_cache.DatasetKey, ['a', 'b'])
     def preprocessing_fn(inputs):
       return {'x': tft.scale_to_0_1(inputs['x'])}
@@ -531,10 +539,7 @@ def mocked_make_cache_entry_key(_):
               specs,
               full_dataset_keys,
               force_tf_compat_v1=use_tf_compat_v1))
-
-    dot_string = nodes.get_dot_graph([analysis_graph_builder._ANALYSIS_GRAPH
-                                     ]).to_string()
-    self.WriteRenderedDotFile(dot_string)
+    self.DebugPublishLatestsRenderedTFTGraph()
     self.assertCountEqual(cache_entry_keys, [mocked_cache_entry_key])
 
   def test_duplicate_label_error(self):
@@ -575,4 +580,4 @@ class _Analyzer(
 
 
 if __name__ == '__main__':
-  test_case.main()
+  tft_unit.main()

tensorflow_transform/beam/beam_nodes.py

@@ -103,6 +103,8 @@ def _get_tensor_type_name(self, tensor):
       return 'Tensor'
     elif isinstance(tensor, tf.SparseTensor):
       return 'SparseTensor'
+    elif isinstance(tensor, tf.RaggedTensor):
+      return 'RaggedTensor'
     raise ValueError('Got a {}, expected a Tensor or SparseTensor'.format(
         type(tensor)))
 

tensorflow_transform/beam/impl_test.py

@@ -30,8 +30,8 @@
 from tensorflow_transform import pretrained_models
 from tensorflow_transform import schema_inference
 import tensorflow_transform.beam as tft_beam
-from tensorflow_transform.beam import tft_unit
 from tensorflow_transform.beam.tft_beam_io import transform_fn_io
+from tensorflow_transform.beam import tft_unit
 from tfx_bsl.tfxio import tensor_adapter
 
 from google.protobuf import text_format

tensorflow_transform/beam/tft_beam_io/beam_metadata_io_test.py

@@ -19,8 +19,8 @@
 import apache_beam as beam
 import tensorflow as tf
 from tensorflow_transform import output_wrapper
-from tensorflow_transform.beam import tft_unit
 from tensorflow_transform.beam.tft_beam_io import beam_metadata_io
+from tensorflow_transform.beam import tft_unit
 from tensorflow_transform.beam.tft_beam_io import test_metadata
 import tensorflow_transform.test_case as tft_test_case
 from tensorflow_transform.tf_metadata import metadata_io

tensorflow_transform/beam/tft_beam_io/transform_fn_io_test.py

@@ -18,11 +18,10 @@
 import apache_beam as beam
 from apache_beam.testing import util as beam_test_util
 import tensorflow as tf
-
 import tensorflow_transform as tft
-from tensorflow_transform.beam import tft_unit
 from tensorflow_transform.beam.tft_beam_io import beam_metadata_io
 from tensorflow_transform.beam.tft_beam_io import transform_fn_io
+from tensorflow_transform.beam import tft_unit
 from tensorflow_transform.beam.tft_beam_io import test_metadata
 from tensorflow_transform.tf_metadata import metadata_io
 

tensorflow_transform/beam/tft_unit.py

@@ -16,12 +16,13 @@
 import os
 import tempfile
 from typing import Dict, Iterable, List, Optional, Tuple
+from absl import logging
 
 import apache_beam as beam
 import pyarrow as pa
 import tensorflow as tf
 import tensorflow_transform as tft
-from tensorflow_transform.beam import impl as beam_impl
+import tensorflow_transform.beam as tft_beam
 from tensorflow_transform.beam.tft_beam_io import transform_fn_io
 from tensorflow_transform import test_case
 from tensorflow_transform.beam import test_helpers
@@ -333,38 +334,47 @@ def assertAnalyzeAndTransformResults(self,
     temp_dir = temp_dir or tempfile.mkdtemp(
         prefix=self._testMethodName, dir=self.get_temp_dir())
     with beam_pipeline or self._makeTestPipeline() as pipeline:
-      with beam_impl.Context(
+      with tft_beam.Context(
           temp_dir=temp_dir,
           desired_batch_size=desired_batch_size,
-          force_tf_compat_v1=force_tf_compat_v1):
+          force_tf_compat_v1=force_tf_compat_v1,
+      ):
         source_ptransform = (
             input_data if isinstance(input_data, beam.PTransform) else
             beam.Create(input_data, reshuffle=False))
         input_data = pipeline | 'CreateInput' >> source_ptransform
         if test_data is None:
           (transformed_data, transformed_metadata), transform_fn = (
-              (input_data, input_metadata)
-              | beam_impl.AnalyzeAndTransformDataset(
-                  preprocessing_fn,
-                  output_record_batches=output_record_batches))
+              input_data,
+              input_metadata,
+          ) | tft_beam.AnalyzeAndTransformDataset(
+              preprocessing_fn, output_record_batches=output_record_batches
+          )
         else:
-          transform_fn = ((input_data, input_metadata)
-                          | beam_impl.AnalyzeDataset(preprocessing_fn))
+          transform_fn = (input_data, input_metadata) | tft_beam.AnalyzeDataset(
+              preprocessing_fn
+          )
           test_data = pipeline | 'CreateTest' >> beam.Create(test_data)
           transformed_data, transformed_metadata = (
-              ((test_data, input_metadata), transform_fn)
-              | beam_impl.TransformDataset(
-                  output_record_batches=output_record_batches))
+              (test_data, input_metadata),
+              transform_fn,
+          ) | tft_beam.TransformDataset(
+              output_record_batches=output_record_batches
+          )
 
         # Write transform_fn so we can test its assets
         _ = transform_fn | transform_fn_io.WriteTransformFn(temp_dir)
 
         transformed_data_path = os.path.join(temp_dir, 'transformed_data')
         if expected_data is not None:
-          _ = ((transformed_data, transformed_metadata)
-               | 'Encode' >> beam_impl.EncodeTransformedDataset()
-               | 'Write' >> beam.io.tfrecordio.WriteToTFRecord(
-                   transformed_data_path, shard_name_template=''))
+          _ = (
+              (transformed_data, transformed_metadata)
+              | 'Encode' >> tft_beam.EncodeTransformedDataset()
+              | 'Write'
+              >> beam.io.tfrecordio.WriteToTFRecord(
+                  transformed_data_path, shard_name_template=''
+              )
+          )
 
     # TODO(ebreck) Log transformed_data somewhere.
     tf_transform_output = tft.TFTransformOutput(temp_dir)
@@ -406,3 +416,18 @@ def assertAnalyzeAndTransformResults(self,
     for filename, file_contents in expected_vocab_file_contents.items():
       full_filename = tf_transform_output.vocabulary_file_by_name(filename)
       self.AssertVocabularyContents(full_filename, file_contents)
+
+  def DebugPublishLatestsRenderedTFTGraph(
+      self, output_file: Optional[str] = None
+  ):
+    """Outputs a rendered graph which may be used for debugging.
+
+    Requires adding the binary resource to the test target:
+    data = ["//third_party/graphviz:dot_binary"]
+
+    Args:
+      output_file: Path to output the rendered graph file.
+    """
+    logging.info(
+        'DebugPublishLatestsRenderedTFTGraph is not currently supported.'
+    )

tensorflow_transform/beam/vocabulary_integration_test.py

@@ -18,13 +18,12 @@
 import os
 
 import apache_beam as beam
-
 import tensorflow as tf
 import tensorflow_transform as tft
 from tensorflow_transform.beam import analyzer_impls
 from tensorflow_transform.beam import impl as beam_impl
-from tensorflow_transform.beam import tft_unit
 from tensorflow_transform.beam.tft_beam_io import transform_fn_io
+from tensorflow_transform.beam import tft_unit
 
 from tensorflow_metadata.proto.v0 import schema_pb2
 

tensorflow_transform/beam/vocabulary_tfrecord_gzip_integration_test.py

@@ -14,8 +14,8 @@
 # limitations under the License.
 """Tests for tfrecord_gzip tft.vocabulary and tft.compute_and_apply_vocabulary."""
 
-from tensorflow_transform.beam import tft_unit
 from tensorflow_transform.beam import vocabulary_integration_test
+from tensorflow_transform.beam import tft_unit
 
 
 class TFRecordVocabularyIntegrationTest(