Adds support for handling struct features in schema_as_feature_spec.

iindyk · tfx-copybara · commit b2ede543d47f · 2022-09-15T12:11:28.000-07:00
Presence of a struct feature is interpreted as physical format being `tf.SequenceExample`. `schema_as_feature_spec` returns a union of context and sequence features in such case. This represents how `preprocessing_fn` sees parsed features (e.g. with TfSequenceExampleRecord tfxio).

PiperOrigin-RevId: 474625238
diff --git a/RELEASE.md b/RELEASE.md
@@ -7,6 +7,8 @@
 *   Introduced `tft.experimental.document_frequency` and `tft.experimental.idf`
     which map each term to its document frequency and inverse document frequency
     in the same order as the terms in documents.
+*   `schema_utils.schema_as_feature_spec` now supports struct features as a way
+    to describe `tf.SequenceExample` data.
 
 ## Bug Fixes and Other Changes
 
diff --git a/tensorflow_transform/tf_metadata/schema_utils.py b/tensorflow_transform/tf_metadata/schema_utils.py
@@ -14,7 +14,7 @@
 """Utilities for using the tf.Metadata Schema within TensorFlow."""
 
 import typing
-from typing import Dict, List, Mapping, Optional, Tuple
+from typing import Dict, List, Mapping, Optional, Tuple, Union
 
 import tensorflow as tf
 
@@ -262,15 +262,19 @@ def schema_as_feature_spec(
   For a Feature with a FixedShape we generate a FixedLenFeature with no default.
   For a Feature without a FixedShape we generate a VarLenFeature.  For a
   SparseFeature we generate a SparseFeature.
+  If schema contains struct feature, then it must also contain
+  TensorRepresentations and is assumed to describe SequenceExample data. The
+  result in such case is union of context and sequence feature specs.
 
   Args:
     schema_proto: A Schema proto.
 
   Returns:
     A pair (feature spec, domains) where feature spec is a dict whose keys are
-        feature names and values are instances of FixedLenFeature, VarLenFeature
-        or SparseFeature, and `domains` is a dict whose keys are feature names
-        and values are one of the `domain_info` oneof, e.g. IntDomain.
+        feature names and values are instances of FixedLenFeature,
+        VarLenFeature, SparseFeature or RaggedFeature, and `domains` is a dict
+        whose keys are feature names and values are one of the `domain_info`
+        oneof, e.g. IntDomain.
 
   Raises:
     ValueError: If the schema proto is invalid.
@@ -285,6 +289,12 @@ def schema_as_feature_spec(
 
   if schema_utils_legacy.get_generate_legacy_feature_spec(schema_proto):
     return _legacy_schema_as_feature_spec(schema_proto)
+
+  # Presence of a struct means that data's physical format is tf.SequenceExample
+  # and the struct contains sequence features.
+  if any(feature.type == schema_pb2.STRUCT for feature in schema_proto.feature):
+    return _sequence_schema_as_feature_spec(schema_proto)
+
   feature_spec = {}
   # Will hold the domain_info (IntDomain, FloatDomain etc.) of the feature.  For
   # sparse features, will hold the domain_info of the values feature.  Features
@@ -335,7 +345,71 @@ def schema_as_feature_spec(
   return SchemaAsFeatureSpecResult(feature_spec, domains)
 
 
-def _get_string_domains(schema):
+def _sequence_schema_as_feature_spec(
+    schema: schema_pb2.Schema) -> SchemaAsFeatureSpecResult:
+  """Generates a feature spec from a Schema describing tf.SequenceExample data.
+
+  See `tensor_representation_util.CreateTfSequenceExampleParserConfig`s
+  docstring for feature spec generation rules.
+  We mix context and sequence feature specs to replicate how preprocessing_fn
+  sees input features -- as top-level values of a single `inputs` dict. Note
+  that this makes the feature spec generation irreversible without additional
+  input since it's no longer possible to distinguish context and sequence
+  features to produce the original schema.
+
+  Args:
+    schema: A TFMD Schema proto.
+
+  Returns:
+    A pair (feature spec, domains) where feature spec is a dict whose keys are
+        feature names and values are instances of FixedLenFeature,
+        VarLenFeature, SparseFeature or RaggedFeature, and `domains` is a dict
+        whose keys are feature names and values are one of the `domain_info`
+        oneof, e.g. IntDomain.
+
+  Raises:
+    ValueError: If `TensorRepresentation`s in the schema result in feature specs
+        that are not supported.
+  """
+  (context_feature_spec, sequence_feature_spec
+  ) = tensor_representation_util.CreateTfSequenceExampleParserConfig(schema)
+  feature_spec = {**context_feature_spec, **sequence_feature_spec}
+  string_domains = _get_string_domains(schema)
+  domain_by_feature_name = _get_source_feature_domains(schema, string_domains)
+  domains = {}
+  for name, spec in feature_spec.items():
+    if isinstance(spec, (tf.io.FixedLenFeature, tf.io.VarLenFeature)):
+      source_feature_name = name
+    elif isinstance(
+        spec, tf.io.SparseFeature) or common_types.is_ragged_feature(spec):
+      source_feature_name = spec.value_key
+    else:
+      raise ValueError('spec is not recognized')
+    if source_feature_name in domain_by_feature_name:
+      domains[name] = domain_by_feature_name[source_feature_name]
+  return SchemaAsFeatureSpecResult(feature_spec, domains)
+
+
+def _get_source_feature_domains(
+    schema_or_domain: Union[schema_pb2.Schema, schema_pb2.StructDomain],
+    string_domains: Dict[str, schema_pb2.StringDomain]
+) -> Dict[str, common_types.DomainType]:
+  """Recursively extracts domains of all source features in the schema."""
+  result = {}
+  for feature in schema_or_domain.feature:
+    domain_info = feature.WhichOneof('domain_info')
+    if domain_info == 'struct_domain':
+      result.update(
+          _get_source_feature_domains(feature.struct_domain, string_domains))
+    else:
+      domain = _get_domain(feature, string_domains)
+      if domain is not None:
+        result[feature.name] = domain
+  return result
+
+
+def _get_string_domains(
+    schema: schema_pb2.Schema) -> Dict[str, schema_pb2.StringDomain]:
   return {domain.name: domain for domain in schema.string_domain}
 
 
diff --git a/tensorflow_transform/tf_metadata/schema_utils_test_cases.py b/tensorflow_transform/tf_metadata/schema_utils_test_cases.py
@@ -787,6 +787,224 @@
       },
   ])
 
+  NON_ROUNDTRIP_SCHEMAS.extend([{
+      'testcase_name':
+          'sequence',
+      'ascii_proto':
+          """
+          feature {
+            name: "int_feature"
+            type: INT
+            value_count {
+              min: 1
+              max: 1
+            }
+          }
+
+          feature {
+            name: "##SEQUENCE##"
+            type: STRUCT
+            struct_domain {
+              feature {
+                name: "int_feature"
+                type: INT
+                value_count {
+                  min: 0
+                  max: 2
+                }
+              }
+            }
+          }
+          tensor_representation_group {
+            key: ""
+            value {
+              tensor_representation {
+                key: "int_feature"
+                value { varlen_sparse_tensor { column_name: "int_feature" } }
+              }
+              tensor_representation {
+                key: "seq_int_feature"
+                value {
+                  ragged_tensor {
+                    feature_path { step: "##SEQUENCE##" step: "int_feature" }
+                  }
+                }
+              }
+            }
+          }
+        """,
+      'feature_spec': {
+          'int_feature':
+              tf.io.VarLenFeature(dtype=tf.int64),
+          'seq_int_feature':
+              tf.io.RaggedFeature(
+                  dtype=tf.int64,
+                  value_key='int_feature',
+                  partitions=[],
+                  row_splits_dtype=tf.int64,
+                  validate=False),
+      },
+  }, {
+      'testcase_name':
+          'sequence_no_context',
+      'ascii_proto':
+          """
+          feature {
+            name: "##SEQUENCE##"
+            type: STRUCT
+            struct_domain {
+              feature {
+                name: "x"
+                type: INT
+                value_count {
+                  min: 0
+                  max: 2
+                }
+              }
+            }
+          }
+          tensor_representation_group {
+            key: ""
+            value {
+              tensor_representation {
+                key: "x"
+                value { ragged_tensor {
+                            feature_path { step: "##SEQUENCE##" step: "x" } } }
+              }
+            }
+          }
+        """,
+      'feature_spec': {
+          'x':
+              tf.io.RaggedFeature(
+                  dtype=tf.int64,
+                  value_key='x',
+                  partitions=[],
+                  row_splits_dtype=tf.int64,
+                  validate=False),
+      },
+  }, {
+      'testcase_name':
+          'sequence_with_domains',
+      'ascii_proto':
+          """
+          feature {
+            name: "int_feature"
+            type: INT
+            value_count {
+              min: 1
+              max: 1
+            }
+            int_domain { min: 0 max: 9 }
+          }
+
+          feature {
+            name: "##SEQUENCE##"
+            type: STRUCT
+            struct_domain {
+              feature {
+                name: "float_feature"
+                type: FLOAT
+                value_count {
+                  min: 0
+                  max: 2
+                }
+                float_domain { min: 1.0}
+              }
+            }
+          }
+          tensor_representation_group {
+            key: ""
+            value {
+              tensor_representation {
+                key: "int_feature"
+                value { varlen_sparse_tensor { column_name: "int_feature" } }
+              }
+              tensor_representation {
+                key: "seq_float_feature"
+                value {
+                  ragged_tensor {
+                    feature_path { step: "##SEQUENCE##" step: "float_feature" }
+                  }
+                }
+              }
+            }
+          }
+        """,
+      'feature_spec': {
+          'int_feature':
+              tf.io.VarLenFeature(dtype=tf.int64),
+          'seq_float_feature':
+              tf.io.RaggedFeature(
+                  dtype=tf.float32,
+                  value_key='float_feature',
+                  partitions=[],
+                  row_splits_dtype=tf.int64,
+                  validate=False),
+      },
+      'domains': {
+          'int_feature': schema_pb2.IntDomain(min=0, max=9),
+          'seq_float_feature': schema_pb2.FloatDomain(min=1.0)
+      }
+  }, {
+      'testcase_name':
+          'sequence_with_string_domain',
+      'ascii_proto':
+          """
+          feature {
+            name: "int_feature"
+            type: INT
+          }
+
+          feature {
+            name: "##SEQUENCE##"
+            type: STRUCT
+            struct_domain {
+              feature {
+                name: "string_feature"
+                type: BYTES
+                value_count {
+                  min: 0
+                  max: 2
+                }
+                string_domain {value: "a" value: "b"}
+              }
+            }
+          }
+          tensor_representation_group {
+            key: ""
+            value {
+              tensor_representation {
+                key: "int_feature"
+                value { varlen_sparse_tensor { column_name: "int_feature" } }
+              }
+              tensor_representation {
+                key: "seq_string_feature"
+                value {
+                  ragged_tensor {
+                    feature_path { step: "##SEQUENCE##" step: "string_feature" }
+                  }
+                }
+              }
+            }
+          }
+        """,
+      'feature_spec': {
+          'int_feature':
+              tf.io.VarLenFeature(dtype=tf.int64),
+          'seq_string_feature':
+              tf.io.RaggedFeature(
+                  dtype=tf.string,
+                  value_key='string_feature',
+                  partitions=[],
+                  row_splits_dtype=tf.int64,
+                  validate=False),
+      },
+      'domains': {
+          'seq_string_feature': schema_pb2.StringDomain(value=['a', 'b'])
+      }
+  }])
+
   INVALID_SCHEMA_PROTOS.extend([{
       'testcase_name':
           'ragged_feature_non_int_row_lengths',
