Adds tft_beam.EncodeTransformedDatasetAsSerializedExamples which can be used to easily encode transformed data in either record batch or instance dict format.

PiperOrigin-RevId: 482514132

Author: zoyahav

RELEASE.md

@@ -17,6 +17,8 @@
 *   TensorRepresentations in schema used for
     `schema_utils.schema_as_feature_spec` can now share name with their source
     features.
+*   Introduced `tft_beam.EncodeTransformedDataset` which can be used to easily
+    encode transformed data in preparation for materialization.
 
 ## Bug Fixes and Other Changes
 

examples/census_example_common.py

@@ -22,7 +22,6 @@
 import tensorflow.compat.v2 as tf
 import tensorflow_transform as tft
 import tensorflow_transform.beam as tft_beam
-from tfx_bsl.public.tfxio import RecordBatchToExamplesEncoder
 from tfx_bsl.public import tfxio
 
 CATEGORICAL_FEATURE_KEYS = [
@@ -210,16 +209,11 @@ def preprocessing_fn(inputs):
           raw_dataset | tft_beam.AnalyzeAndTransformDataset(
               preprocessing_fn, output_record_batches=True))
 
-      # Transformed metadata is not necessary for encoding.
-      transformed_data, _ = transformed_dataset
-
       # Extract transformed RecordBatches, encode and write them to the given
       # directory.
-      coder = RecordBatchToExamplesEncoder()
       _ = (
-          transformed_data
-          | 'EncodeTrainData' >>
-          beam.FlatMapTuple(lambda batch, _: coder.encode(batch))
+          transformed_dataset
+          | 'EncodeTrainData' >> tft_beam.EncodeTransformedDataset()
           | 'WriteTrainData' >> beam.io.WriteToTFRecord(
               os.path.join(working_dir, TRANSFORMED_TRAIN_DATA_FILEBASE)))
 
@@ -243,15 +237,11 @@ def preprocessing_fn(inputs):
           (raw_test_dataset, transform_fn)
           | tft_beam.TransformDataset(output_record_batches=True))
 
-      # Transformed metadata is not necessary for encoding.
-      transformed_test_data, _ = transformed_test_dataset
-
       # Extract transformed RecordBatches, encode and write them to the given
       # directory.
       _ = (
-          transformed_test_data
-          | 'EncodeTestData' >>
-          beam.FlatMapTuple(lambda batch, _: coder.encode(batch))
+          transformed_test_dataset
+          | 'EncodeTestData' >> tft_beam.EncodeTransformedDataset()
           | 'WriteTestData' >> beam.io.WriteToTFRecord(
               os.path.join(working_dir, TRANSFORMED_TEST_DATA_FILEBASE)))
 

examples/sentiment_example.py

@@ -189,32 +189,28 @@ def preprocessing_fn(inputs):
             LABEL_KEY: inputs[LABEL_KEY]
         }
 
-      # Transformed metadata is not necessary for encoding.
       # The TFXIO output format is chosen for improved performance.
-      (transformed_train_data, _), transform_fn = (
+      transformed_train_data, transform_fn = (
           (train_data, tfxio_train_data.TensorAdapterConfig())
           | 'AnalyzeAndTransform' >> tft_beam.AnalyzeAndTransformDataset(
               preprocessing_fn, output_record_batches=True))
 
-      transformed_test_data, _ = (
+      transformed_test_data = (
           ((test_data, tfxio_test_data.TensorAdapterConfig()), transform_fn)
           |
           'Transform' >> tft_beam.TransformDataset(output_record_batches=True))
 
       # Extract transformed RecordBatches, encode and write them to the given
       # directory.
-      coder = tfxio.RecordBatchToExamplesEncoder()
       _ = (
           transformed_train_data
-          | 'EncodeTrainData' >>
-          beam.FlatMapTuple(lambda batch, _: coder.encode(batch))
+          | 'EncodeTrainData' >> tft_beam.EncodeTransformedDataset()
           | 'WriteTrainData' >> beam.io.WriteToTFRecord(
               os.path.join(working_dir, TRANSFORMED_TRAIN_DATA_FILEBASE)))
 
       _ = (
           transformed_test_data
-          | 'EncodeTestData' >>
-          beam.FlatMapTuple(lambda batch, _: coder.encode(batch))
+          | 'EncodeTestData' >> tft_beam.EncodeTransformedDataset()
           | 'WriteTestData' >> beam.io.WriteToTFRecord(
               os.path.join(working_dir, TRANSFORMED_TEST_DATA_FILEBASE)))
 

tensorflow_transform/beam/__init__.py

@@ -24,6 +24,7 @@
 from tensorflow_transform.beam.impl import AnalyzeAndTransformDataset
 from tensorflow_transform.beam.impl import AnalyzeDataset
 from tensorflow_transform.beam.impl import AnalyzeDatasetWithCache
+from tensorflow_transform.beam.impl import EncodeTransformedDataset
 from tensorflow_transform.beam.impl import TransformDataset
 from tensorflow_transform.beam.tft_beam_io import *
 

tensorflow_transform/beam/impl.py

@@ -79,6 +79,7 @@
 from tensorflow_transform.tf_metadata import dataset_metadata
 from tensorflow_transform.tf_metadata import metadata_io
 from tensorflow_transform.tf_metadata import schema_utils
+from tfx_bsl.coders import example_coder
 from tfx_bsl.telemetry import collection as telemetry
 from tfx_bsl.telemetry import util as telemetry_util
 from tfx_bsl.tfxio import tensor_representation_util
@@ -1499,10 +1500,13 @@ def expand(self, dataset_and_transform_fn):
       deferred_schema = (
           output_metadata.deferred_metadata
           | 'GetDeferredSchema' >> beam.Map(lambda m: m.schema))
+      output_dataset_metadata = output_metadata.dataset_metadata
     else:
       deferred_schema = (
           self.pipeline
           | 'CreateDeferredSchema' >> beam.Create([output_metadata.schema]))
+      output_dataset_metadata = output_metadata
+    output_dataset_metadata._output_record_batches = self._output_record_batches  # pylint: disable=protected-access
 
     # Increment input metrics.
     _ = (
@@ -1573,3 +1577,67 @@ def expand(self, dataset_and_transform_fn):
     _clear_shared_state_after_barrier(self.pipeline, output_data)
 
     return (output_data, output_metadata)
+
+
+class EncodeTransformedDataset(beam.PTransform):
+  """Encodes transformed data into serialized tf.Examples.
+
+  Should operate on the output of `TransformDataset`, this can operate on either
+  record batch or instance dict data.
+  The expected input is a (transformed_data, transformed_metadata) tuple.
+
+  Example use:
+
+  >>> def preprocessing_fn(inputs):
+  ...   return {'x_scaled': tft.scale_to_z_score(inputs['x'], name='x')}
+  >>> raw_data = [dict(x=1), dict(x=2), dict(x=3)]
+  >>> feature_spec = dict(x=tf.io.FixedLenFeature([], tf.int64))
+  >>> raw_data_metadata = tft.DatasetMetadata.from_feature_spec(feature_spec)
+  >>> output_path = os.path.join(tempfile.mkdtemp(), 'result')
+  >>> with beam.Pipeline() as p:
+  ...   with tft_beam.Context(temp_dir=tempfile.mkdtemp()):
+  ...     data_pcoll = p | beam.Create(raw_data)
+  ...     transformed_dataset, transform_fn = (
+  ...         (data_pcoll, raw_data_metadata)
+  ...         | tft_beam.AnalyzeAndTransformDataset(preprocessing_fn))
+  ...     _ = (
+  ...       transformed_dataset
+  ...       | tft_beam.EncodeTransformedDataset()
+  ...       | beam.io.WriteToTFRecord(output_path, shard_name_template=''))
+  >>> result_feature_spec ={'x_scaled': tf.io.FixedLenFeature([], tf.float32)}
+  >>> list(tf.data.TFRecordDataset([output_path])
+  ...      .map(lambda x: tf.io.parse_example(x, result_feature_spec))
+  ...      .as_numpy_iterator())
+  [{'x_scaled': -1.2247448}, {'x_scaled': 0.0}, {'x_scaled': 1.2247448}]
+  """
+
+  def _extract_input_pvalues(self, transformed_data_and_metadata):
+    # This method lets beam know that metadata is not a pvalue.
+    return transformed_data_and_metadata, [transformed_data_and_metadata[0]]
+
+  def expand(self, transformed_data_and_metadata):
+
+    transformed_data, transformed_metadata = transformed_data_and_metadata
+
+    deferred_schema = (
+        transformed_metadata.deferred_metadata
+        | 'GetDeferredSchema' >> beam.Map(lambda m: m.schema))
+
+    if transformed_metadata.dataset_metadata._output_record_batches:  # pylint: disable=protected-access
+      transformed_data_coder_pcol = (
+          deferred_schema | 'RecordBatchToExamplesEncoder' >> beam.Map(
+              example_coder.RecordBatchToExamplesEncoder))
+      encode_ptransform = 'EncodeRecordBatches' >> beam.FlatMap(
+          # Dropping passthrough features.
+          lambda elem, coder: coder.encode(elem[0]),
+          coder=beam.pvalue.AsSingleton(transformed_data_coder_pcol))
+    else:
+      transformed_data_coder_pcol = (
+          deferred_schema
+          | 'ExampleProtoCoder' >> beam.Map(
+              example_proto_coder.ExampleProtoCoder))
+      encode_ptransform = 'EncodeInstances' >> beam.Map(
+          lambda data, data_coder: data_coder.encode(data),
+          data_coder=beam.pvalue.AsSingleton(transformed_data_coder_pcol))
+
+    return transformed_data | encode_ptransform

tensorflow_transform/beam/impl_test.py

@@ -34,7 +34,6 @@
 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 tfx_bsl.coders import example_coder
 from tfx_bsl.tfxio import tensor_adapter
 
 from google.protobuf import text_format
@@ -4418,27 +4417,11 @@ def test3dSparseWithTFXIO(self):
             | tft_beam.AnalyzeAndTransformDataset(
                 lambda inputs: inputs,
                 output_record_batches=self._OutputRecordBatches()))
-        if self._OutputRecordBatches():
 
-          def record_batch_to_examples(data_batch):
-            # Ignore unary pass-through features.
-            record_batch, _ = data_batch
-            return example_coder.RecordBatchToExamples(record_batch)
-
-          transformed_and_serialized = (
-              transformed_data |
-              'EncodeTransformedData' >> beam.FlatMap(record_batch_to_examples))
-        else:
-          transformed_data_coder = tft.coders.ExampleProtoCoder(
-              transformed_metadata.schema)
-          transformed_and_serialized = (
-              transformed_data | 'EncodeTransformedData' >> beam.Map(
-                  transformed_data_coder.encode))
-
-        _ = (
-            transformed_and_serialized
-            | 'Write' >> beam.io.WriteToTFRecord(
-                materialize_path, shard_name_template=''))
+        _ = ((transformed_data, transformed_metadata)
+             | 'Encode' >> tft_beam.EncodeTransformedDataset()
+             | 'Write' >> beam.io.WriteToTFRecord(
+                 materialize_path, shard_name_template=''))
         _ = (
             transform_fn
             | 'WriteTransformFn' >>

tensorflow_transform/beam/tft_unit.py

@@ -22,7 +22,6 @@
 import tensorflow as tf
 import tensorflow_transform as tft
 from tensorflow_transform.beam import impl as beam_impl
-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 import test_case
 from tensorflow_transform.beam import test_helpers
@@ -230,7 +229,7 @@ def preprocessing_fn(inputs):
         try:
           output_tensor.set_shape(output_shape)
         except ValueError as e:
-          raise ValueError('Error for key {}: {}'.format(key, str(e)))
+          raise ValueError(f'Error for key {key}') from e
         # Add a batch dimension
         output_tensor = tf.expand_dims(output_tensor, 0)
         # Broadcast along the batch dimension
@@ -362,41 +361,10 @@ def assertAnalyzeAndTransformResults(self,
 
         transformed_data_path = os.path.join(temp_dir, 'transformed_data')
         if expected_data is not None:
-          if isinstance(transformed_metadata,
-                        beam_metadata_io.BeamDatasetMetadata):
-            deferred_schema = (
-                transformed_metadata.deferred_metadata
-                | 'GetDeferredSchema' >> beam.Map(lambda m: m.schema))
-          else:
-            deferred_schema = (
-                self.pipeline | 'CreateDeferredSchema' >> beam.Create(
-                    [transformed_metadata.schema]))
-
-          if output_record_batches:
-            # Since we are using a deferred schema, obtain a pcollection
-            # containing the data coder that will be created from it.
-            transformed_data_coder_pcol = (
-                deferred_schema | 'RecordBatchToExamplesEncoder' >> beam.Map(
-                    example_coder.RecordBatchToExamplesEncoder))
-
-            encode_ptransform = 'EncodeRecordBatches' >> beam.FlatMap(
-                _encode_transformed_data_batch,
-                coder=beam.pvalue.AsSingleton(transformed_data_coder_pcol))
-          else:
-            # Since we are using a deferred schema, obtain a pcollection
-            # containing the data coder that will be created from it.
-            transformed_data_coder_pcol = (
-                deferred_schema
-                | 'ExampleProtoCoder' >> beam.Map(tft.coders.ExampleProtoCoder))
-            encode_ptransform = 'EncodeExamples' >> beam.Map(
-                lambda data, data_coder: data_coder.encode(data),
-                data_coder=beam.pvalue.AsSingleton(transformed_data_coder_pcol))
-
-          _ = (
-              transformed_data
-              | encode_ptransform
-              | beam.io.tfrecordio.WriteToTFRecord(
-                  transformed_data_path, shard_name_template=''))
+          _ = ((transformed_data, transformed_metadata)
+               | 'Encode' >> beam_impl.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)

tensorflow_transform/tf_metadata/dataset_metadata.py

@@ -37,6 +37,7 @@ class DatasetMetadata:
 
   def __init__(self, schema: schema_pb2.Schema):
     self._schema = schema
+    self._output_record_batches = True
 
   @classmethod
   def from_feature_spec(
@@ -53,7 +54,7 @@ def schema(self) -> schema_pb2.Schema:
 
   def __eq__(self, other):
     if isinstance(other, self.__class__):
-      return self.__dict__ == other.__dict__
+      return self.schema == other.schema
     return NotImplemented
 
   def __ne__(self, other):