Automated rollback of commit 067cf11

PiperOrigin-RevId: 551536436

Author: zoyahav

tensorflow_transform/beam/impl.py

@@ -1360,11 +1360,17 @@ def expand(self, dataset):
 
 def _remove_columns_from_metadata(metadata, excluded_columns):
   """Remove columns from metadata without mutating original metadata."""
-  feature_spec, domains = schema_utils.schema_as_feature_spec(metadata.schema)
-  new_feature_spec = {name: spec for name, spec in feature_spec.items()
-                      if name not in excluded_columns}
-  new_domains = {name: spec for name, spec in domains.items()
-                 if name not in excluded_columns}
+  generated = schema_utils.schema_as_feature_spec(metadata.schema)
+  new_feature_spec = {
+      name: spec
+      for name, spec in generated.feature_spec.items()
+      if name not in excluded_columns
+  }
+  new_domains = {
+      name: spec
+      for name, spec in generated.domains.items()
+      if name not in excluded_columns
+  }
   return dataset_metadata.DatasetMetadata.from_feature_spec(
       new_feature_spec, new_domains)
 

tensorflow_transform/tf_metadata/schema_utils.py

@@ -13,17 +13,14 @@
 # limitations under the License.
 """Utilities for using the tf.Metadata Schema within TensorFlow."""
 
+import dataclasses
 import typing
 from typing import Dict, List, Mapping, Optional, Tuple, Union
 
 import tensorflow as tf
-
 from tensorflow_transform import common_types
 from tensorflow_transform.tf_metadata import schema_utils_legacy
 from tfx_bsl.tfxio import tensor_representation_util
-# TODO(b/243513856): Switch to `collections.namedtuple` or `typing.NamedTuple`
-# once the Spark issue is resolved.
-from tfx_bsl.types import tfx_namedtuple
 
 from tensorflow_metadata.proto.v0 import path_pb2
 from tensorflow_metadata.proto.v0 import schema_pb2
@@ -245,10 +242,15 @@ def _set_domain(name, feature, domain):
     raise ValueError('Feature "{}" has invalid domain {}'.format(name, domain))
 
 
-SchemaAsFeatureSpecResult = tfx_namedtuple.TypedNamedTuple(
-    'SchemaAsFeatureSpecResult',
-    [('feature_spec', Dict[str, common_types.FeatureSpecType]),
-     ('domains', Dict[str, common_types.DomainType])])
+@dataclasses.dataclass(frozen=True)
+class SchemaAsFeatureSpecResult:
+  feature_spec: Dict[str, common_types.FeatureSpecType]
+  domains: Dict[str, common_types.DomainType]
+
+  # This is needed because many users unpack this with:
+  # `feature_spec, domains = schema_utils.schema_as_feature_spec()`.
+  def __iter__(self):
+    return (getattr(self, field.name) for field in dataclasses.fields(self))
 
 
 def _standardize_default_value(

tensorflow_transform/tf_metadata/schema_utils_test.py

@@ -46,7 +46,10 @@ def test_schema_as_feature_spec(
     schema_utils_legacy.set_generate_legacy_feature_spec(
         schema_proto, generate_legacy_feature_spec)
     result = schema_utils.schema_as_feature_spec(schema_proto)
-    self.assertEqual(result, (feature_spec, domains or {}))
+    self.assertEqual(
+        result,
+        schema_utils.SchemaAsFeatureSpecResult(feature_spec, domains or {}),
+    )
 
   @parameterized.named_parameters(
       *schema_utils_test_cases.INVALID_SCHEMA_PROTOS)