Experimental support for passing a schema to TFXIO for StatsGen.

PiperOrigin-RevId: 492221403

Author: zwestrick

tensorflow_data_validation/__init__.py

@@ -56,6 +56,7 @@
 
 
 # Import schema utilities.
+from tensorflow_data_validation.utils.schema_util import generate_dummy_schema_with_paths
 from tensorflow_data_validation.utils.schema_util import get_domain
 from tensorflow_data_validation.utils.schema_util import get_feature
 from tensorflow_data_validation.utils.schema_util import load_schema_text

tensorflow_data_validation/statistics/stats_impl.py

@@ -16,7 +16,7 @@
 
 import math
 import random
-from typing import Any, Callable, cast, Dict, Iterable, List, Optional, Text, Tuple
+from typing import Any, Callable, cast, Dict, Iterable, List, Optional, Text, Tuple, Union
 
 import apache_beam as beam
 import pyarrow as pa
@@ -465,7 +465,8 @@ def _schema_has_natural_language_domains(schema: schema_pb2.Schema) -> bool:
 
 def _filter_features(
     record_batch: pa.RecordBatch,
-    feature_allowlist: List[types.FeatureName]) -> pa.RecordBatch:
+    feature_allowlist: Union[List[types.FeatureName], List[types.FeaturePath]]
+) -> pa.RecordBatch:
   """Removes features that are not on the allowlist.
 
   Args:
@@ -478,6 +479,9 @@ def _filter_features(
   columns_to_select = []
   column_names_to_select = []
   for feature_name in feature_allowlist:
+    if isinstance(feature_name, types.FeaturePath):
+      # TODO(b/255895499): Support paths.
+      raise NotImplementedError
     col = arrow_util.get_column(record_batch, feature_name, missing_ok=True)
     if col is None:
       continue

tensorflow_data_validation/statistics/stats_options.py

@@ -22,7 +22,7 @@
 import json
 import logging
 import types as python_types
-from typing import Dict, List, Optional, Text
+from typing import Dict, List, Optional, Text, Union
 
 from tensorflow_data_validation import types
 from tensorflow_data_validation.statistics.generators import stats_generator
@@ -76,14 +76,17 @@ def __init__(
                                                  types.FeatureName]] = None,
       vocab_paths: Optional[Dict[types.VocabName, types.VocabPath]] = None,
       add_default_generators: bool = True,
-      feature_allowlist: Optional[List[types.FeatureName]] = None,
+      # TODO(b/255895499): Support "from schema" for feature_allowlist.
+      feature_allowlist: Optional[Union[List[types.FeatureName],
+                                        List[types.FeaturePath]]] = None,
       experimental_use_sketch_based_topk_uniques: Optional[bool] = None,
       use_sketch_based_topk_uniques: Optional[bool] = None,
       experimental_slice_functions: Optional[List[types.SliceFunction]] = None,
       experimental_slice_sqls: Optional[List[Text]] = None,
       experimental_result_partitions: int = 1,
       experimental_num_feature_partitions: int = 1,
-      slicing_config: Optional[slicing_spec_pb2.SlicingConfig] = None):
+      slicing_config: Optional[slicing_spec_pb2.SlicingConfig] = None,
+      experimental_filter_read_paths: bool = False):
     """Initializes statistics options.
 
     Args:
@@ -151,7 +154,7 @@ def __init__(
         (controlled by `enable_semantic_domain_stats`) and 4) schema-based
         generators that are enabled based on information provided in the schema.
       feature_allowlist: An optional list of names of the features to calculate
-        statistics for.
+        statistics for, or a list of paths.
       experimental_use_sketch_based_topk_uniques: Deprecated, prefer
         use_sketch_based_topk_uniques.
       use_sketch_based_topk_uniques: if True, use the sketch based
@@ -193,8 +196,11 @@ def __init__(
         number of features in a dataset, and never more than the available beam
         parallelism.
       slicing_config: an optional SlicingConfig. SlicingConfig includes
-      slicing_specs specified with feature keys, feature values or slicing
-      SQL queries.
+        slicing_specs specified with feature keys, feature values or slicing
+        SQL queries.
+      experimental_filter_read_paths: If provided, tries to push down either
+        paths passed via feature_allowlist or via the schema (in that priority)
+        to the underlying read operation. Support depends on the file reader.
     """
     self.generators = generators
     self.feature_allowlist = feature_allowlist
@@ -241,6 +247,7 @@ def __init__(
     self.experimental_num_feature_partitions = experimental_num_feature_partitions
     self.experimental_result_partitions = experimental_result_partitions
     self.slicing_config = slicing_config
+    self.experimental_filter_read_paths = experimental_filter_read_paths
 
   def to_json(self) -> Text:
     """Convert from an object to JSON representation of the __dict__ attribute.
@@ -340,12 +347,16 @@ def generators(
     self._generators = generators
 
   @property
-  def feature_allowlist(self) -> Optional[List[types.FeatureName]]:
+  def feature_allowlist(
+      self
+  ) -> Optional[Union[List[types.FeatureName], List[types.FeaturePath]]]:
     return self._feature_allowlist
 
   @feature_allowlist.setter
   def feature_allowlist(
-      self, feature_allowlist: Optional[List[types.FeatureName]]) -> None:
+      self, feature_allowlist: Optional[Union[List[types.FeatureName],
+                                              List[types.FeaturePath]]]
+  ) -> None:
     if feature_allowlist is not None and not isinstance(feature_allowlist,
                                                         list):
       raise TypeError('feature_allowlist is of type %s, should be a list.' %
@@ -554,6 +565,14 @@ def experimental_num_feature_partitions(self,
       raise ValueError('experimental_num_feature_partitions must be > 0.')
     self._experimental_num_feature_partitions = feature_partitions
 
+  @property
+  def experimental_filter_read_paths(self) -> bool:
+    return self._experimental_filter_read_paths
+
+  @experimental_filter_read_paths.setter
+  def experimental_filter_read_paths(self, filter_read: bool) -> None:
+    self._experimental_filter_read_paths = filter_read
+
 
 def _validate_sql(sql_query: Text, schema: schema_pb2.Schema):
   arrow_schema = example_coder.ExamplesToRecordBatchDecoder(

tensorflow_data_validation/statistics/stats_options_test.py

@@ -474,7 +474,8 @@ def test_stats_options_from_json(self,
       "_slice_sqls": null,
       "_experimental_result_partitions": 1,
       "_experimental_num_feature_partitions": 1,
-      "_slicing_config": null
+      "_slicing_config": null,
+      "_experimental_filter_read_paths": false
     """
     options_json += type_name_line + '}'
     if want_exception:

tensorflow_data_validation/utils/schema_util.py

@@ -374,3 +374,37 @@ def _recursion_helper(
   result = []
   _recursion_helper(types.FeaturePath([]), schema.feature, result)
   return result
+
+
+def _paths_to_tree(paths: List[types.FeaturePath]):
+  """Convert paths to recursively nested dict."""
+  nested_dict = lambda: collections.defaultdict(nested_dict)
+
+  result = nested_dict()
+
+  def _add(tree, path):
+    if not path:
+      return
+    children = tree[path[0]]
+    _add(children, path[1:])
+
+  for path in paths:
+    _add(result, path.steps())
+  return result
+
+
+def generate_dummy_schema_with_paths(
+    paths: List[types.FeaturePath]) -> schema_pb2.Schema:
+  """Generate a schema with the requested paths and no other information."""
+  schema = schema_pb2.Schema()
+  tree = _paths_to_tree(paths)
+
+  def _add(container, name, children):
+    container.feature.add(name=name)
+    if children:
+      for child_name, grandchildren in children.items():
+        _add(container.feature[-1].struct_domain, child_name, grandchildren)
+
+  for name, children in tree.items():
+    _add(schema, name, children)
+  return schema

tensorflow_data_validation/utils/schema_util_test.py

@@ -754,5 +754,35 @@ def test_look_up_feature(self):
         schema_util.look_up_feature('feature2', container), feature_2)
     self.assertIsNone(schema_util.look_up_feature('feature3', container), None)
 
+  def test_generate_dummy_schema_with_paths(self):
+    schema = text_format.Parse(
+        """
+    feature {
+      name: "foo"
+    }
+    feature {
+      name: "bar"
+    }
+    feature {
+      name: "baz"
+      struct_domain: {
+        feature {
+          name: "zip"
+        }
+        feature {
+          name: "zop"
+        }
+      }
+    }
+    """, schema_pb2.Schema())
+    self.assertEqual(
+        schema_util.generate_dummy_schema_with_paths([
+            types.FeaturePath(['foo']),
+            types.FeaturePath(['bar']),
+            types.FeaturePath(['baz', 'zip']),
+            types.FeaturePath(['baz', 'zop'])
+        ]), schema)
+
+
 if __name__ == '__main__':
   absltest.main()